Performance Improvement of create login in BBF (#3291)

Engine PR : babelfish-for-postgresql/postgresql_modified_for_babelfish#507 In Major version 16 Babelfish introduced a new role bbf_admin_role which is a member of all the roles which created implicitly from TSQL. This increases the Load on any Grant/Revoke operations that happen since in BBF to "select the role through which it gets permission to current user to grant", is inefficient and loops over members of members to find the all members of the role even though the required role is find. This commit optimizes the selection of role with is_admin_option true for babelfish by escaping the search for any TSQL roles conditionally. Signed-off-by: ANJU BHARTI <[email protected]> Issues Resolved BABEL-5349
babelfish-for-postgresql · Jan 7, 2025 · ec665c4 · ec665c4
1 parent be6b58f
commit ec665c4
Show file tree

Hide file tree

Showing 21 changed files with 1,284 additions and 8 deletions.
diff --git a/contrib/babelfishpg_tsql/src/catalog.c b/contrib/babelfishpg_tsql/src/catalog.c
@@ -9,6 +9,7 @@
 #include "catalog/catalog.h"
 #include "catalog/heap.h"
 #include "catalog/indexing.h"
+#include "catalog/pg_auth_members.h"
 #include "catalog/pg_namespace.h"
 #include "catalog/pg_authid.h"
 #include "catalog/pg_proc.h"
@@ -6287,3 +6288,49 @@ get_tvp_typename_typeschemaname(char *proc_name, char *target_arg_name, char **t
 	if(!xactStarted)
 		CommitTransactionCommand();
 }
+
+/*
+ * Checks if the given member is 'bbf_role_admin' role and has the privilege 
+ * to grant a specified role through direct membership. This privilege is 
+ * determined by the presence of the 'admin_option' flag, which allows the 
+ * member to grant the role to others. 
+ * 
+ * @return true if the member is 'bbf_role_admin' and also has the admin option 
+ *         to grant the specified role; otherwise, returns false.
+ */
+bool
+bbf_check_member_has_direct_priv_to_grant_role(Oid member, Oid role)
+{
+	CatCList   *memlist;
+
+	/* 
+	 * TSQL specific behavior and member must be bbf_role_admin.
+	 * If not return false. 
+	 */
+	if (sql_dialect != SQL_DIALECT_TSQL || !IS_TDS_CONN()
+		|| member !=  get_bbf_role_admin_oid())
+		return false;
+
+	/* 
+	 * Find all existing tuples for given role
+	 * whose member is bbf_role_admin
+	 */
+	memlist = SearchSysCacheList2(AUTHMEMROLEMEM,
+									ObjectIdGetDatum(role),
+									ObjectIdGetDatum(member));
+	for (int i = 0; i < memlist->n_members; i++)
+	{
+		HeapTuple	tup = &memlist->members[i]->tuple;
+		Form_pg_auth_members form = (Form_pg_auth_members) GETSTRUCT(tup);
+
+		/* Return true if admin_option is true */
+		if (form->admin_option)
+		{
+			ReleaseSysCacheList(memlist);
+			return true;
+		}
+	}
+
+	ReleaseSysCacheList(memlist);
+	return false;
+}
diff --git a/contrib/babelfishpg_tsql/src/catalog.h b/contrib/babelfishpg_tsql/src/catalog.h
@@ -28,6 +28,7 @@ extern bool IsPLtsqlExtendedCatalog(Oid relationId);
 extern bool IsPltsqlToastRelationHook(Relation relation);
 extern bool IsPltsqlToastClassHook(Form_pg_class pg_class_tup);
 extern void pltsql_drop_relation_refcnt_hook(Relation relation);
+extern bool bbf_check_member_has_direct_priv_to_grant_role(Oid member, Oid role);
 
 /*****************************************
  *			SYS schema

diff --git a/contrib/babelfishpg_tsql/src/hooks.c b/contrib/babelfishpg_tsql/src/hooks.c
@@ -302,6 +302,7 @@ static pltsql_replace_non_determinstic_hook_type prev_pltsql_replace_non_determi
 static pltsql_is_partitioned_table_reloptions_allowed_hook_type prev_pltsql_is_partitioned_table_reloptions_allowed_hook = NULL;
 static ExecFuncProc_AclCheck_hook_type prev_ExecFuncProc_AclCheck_hook = NULL;
 static bbf_execute_grantstmt_as_dbsecadmin_hook_type prev_bbf_execute_grantstmt_as_dbsecadmin_hook = NULL;
+static bbf_check_member_has_direct_priv_to_grant_role_hook_type prev_bbf_check_member_has_direct_priv_to_grant_role_hook = NULL;
 
 /*****************************************
  * 			Install / Uninstall
@@ -517,6 +518,9 @@ InstallExtendedHooks(void)
 
 	prev_bbf_execute_grantstmt_as_dbsecadmin_hook = bbf_execute_grantstmt_as_dbsecadmin_hook;
 	bbf_execute_grantstmt_as_dbsecadmin_hook = handle_grantstmt_for_dbsecadmin;
+
+	prev_bbf_check_member_has_direct_priv_to_grant_role_hook = bbf_check_member_has_direct_priv_to_grant_role_hook;
+	bbf_check_member_has_direct_priv_to_grant_role_hook = bbf_check_member_has_direct_priv_to_grant_role;
 
 	pltsql_get_object_identity_event_trigger_hook = pltsql_get_object_identity_event_trigger;
 
@@ -594,6 +598,8 @@ UninstallExtendedHooks(void)
 	pltsql_is_partitioned_table_reloptions_allowed_hook = prev_pltsql_is_partitioned_table_reloptions_allowed_hook;	
 	ExecFuncProc_AclCheck_hook = prev_ExecFuncProc_AclCheck_hook;
 	bbf_execute_grantstmt_as_dbsecadmin_hook = prev_bbf_execute_grantstmt_as_dbsecadmin_hook;
+	bbf_check_member_has_direct_priv_to_grant_role_hook = prev_bbf_check_member_has_direct_priv_to_grant_role_hook;
+
 
 	bbf_InitializeParallelDSM_hook = NULL;
 	bbf_ParallelWorkerMain_hook = NULL;

diff --git a/test/JDBC/expected/db_securityadmin-vu-verify.out b/test/JDBC/expected/db_securityadmin-vu-verify.out
@@ -1312,6 +1312,15 @@ t
 ~~END~~
 
 
+-- Wait to sync with another session
+SELECT pg_sleep(1);
+go
+~~START~~
+void
+
+~~END~~
+
+
 -- tsql
 DROP LOGIN db_securityadmin_restrictions_login;
 GO
diff --git a/test/JDBC/expected/test_windows_login-vu-prepare.out b/test/JDBC/expected/test_windows_login-vu-prepare.out
@@ -1,5 +1,4 @@
 -- psql
--- TODO: BABEL-5349
 do 
 $$ begin 
     if not exists (select * from pg_catalog.pg_roles where rolname = 'rds_ad') 

diff --git a/test/JDBC/expected/test_windows_login-vu-verify.out b/test/JDBC/expected/test_windows_login-vu-verify.out
@@ -208,3 +208,27 @@ int8
 1
 ~~END~~
 
+
+-- validates bbf_role_admin is direct member of all the role entries of babelfish_authid_user_ext and babelfish_authid_login_ext with admin option
+-- show entries of which bbf_role_admin is not direct member of.
+-- should be 0
+WITH bbf_admin_roles AS (
+    SELECT r.rolname
+    FROM pg_auth_members m
+    JOIN pg_roles r ON (m.roleid = r.oid)
+    WHERE m.member IN (SELECT oid FROM pg_roles WHERE rolname = 'bbf_role_admin')
+    AND m.admin_option = true
+)
+SELECT COUNT(*)
+FROM (
+    SELECT rolname FROM sys.babelfish_authid_login_ext WHERE rolname != 'bbf_role_admin'
+    UNION ALL
+    SELECT rolname FROM sys.babelfish_authid_user_ext
+)
+WHERE rolname NOT IN (SELECT rolname FROM bbf_admin_roles);
+GO
+~~START~~
+int8
+0
+~~END~~
+
diff --git a/test/JDBC/expected/test_windows_login_before-17_2-vu-cleanup.out b/test/JDBC/expected/test_windows_login_before-17_2-vu-cleanup.out
@@ -0,0 +1,91 @@
+-- tsql
+DROP LOGIN PassWordUser;
+GO
+
+-- test for drop login with upn format
+DROP LOGIN [aduser@AD];
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the login 'aduser@ad', because it does not exist or you do not have permission.)~~
+
+
+DROP LOGIN [ad\adUSer];
+GO
+
+DROP LOGIN [ad\aduserDB];
+GO
+
+DROP LOGIN [ad\Aduserlanguage];
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the login 'aduserlanguage@AD', because it does not exist or you do not have permission.)~~
+
+
+DROP LOGIN [ad\Aduserdblanguage];
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the login 'aduserdblanguage@AD', because it does not exist or you do not have permission.)~~
+
+
+DROP LOGIN [ad\AduserdbEnglish];
+GO
+
+DROP LOGIN [ad\Aduseroption];
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the login 'aduseroption@AD', because it does not exist or you do not have permission.)~~
+
+
+DROP LOGIN [babel\aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa]
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the login 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa@BABEL', because it does not exist or you do not have permission.)~~
+
+
+DROP LOGIN [ba.bel\username];
+GO
+
+-- test for non-existent login
+DROP LOGIN [babel\non_existent_login]
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: Cannot drop the login 'non_existent_login@BABEL', because it does not exist or you do not have permission.)~~
+
+
+-- drop login with different casing
+DROP LOGIN UserPassword;
+GO
+
+DROP LOGIN [BabeL\DupLicate];
+GO
+
+DROP LOGIN [Babel\DuplicateDefaultDB];
+GO
+
+-- test drop logins for logins with different language names
+-- Arabic
+DROP LOGIN [babel\كلب];
+GO
+
+-- Mongolian
+DROP LOGIN [babel\өглөө];
+GO
+
+-- Greek
+DROP LOGIN [babel\ελπίδα];
+GO
+
+-- Chinese
+DROP LOGIN [babel\爱];
+GO
+
+DROP DATABASE ad_db;
+GO
+
+