postgresql: Move PgMembership to module_utils (#61598)

* postgresql: move PgMembership to module_utils

* postgresql: move PgMembership to module_utils, optimize logic

lib/ansible/module_utils/postgres.py

@@ -35,6 +35,7 @@ except ImportError:
     HAS_PSYCOPG2 = False
 
 from ansible.module_utils.basic import missing_required_lib
+from ansible.module_utils.database import pg_quote_identifier
 from ansible.module_utils._text import to_native
 from ansible.module_utils.six import iteritems
 from distutils.version import LooseVersion
@@ -197,3 +198,113 @@ def get_conn_params(module, params_dict, warn_db_default=True):
         kw["host"] = params_dict["login_unix_socket"]
 
     return kw
+
+
+class PgMembership(object):
+    def __init__(self, module, cursor, groups, target_roles, fail_on_role=True):
+        self.module = module
+        self.cursor = cursor
+        self.target_roles = [r.strip() for r in target_roles]
+        self.groups = [r.strip() for r in groups]
+        self.executed_queries = []
+        self.granted = {}
+        self.revoked = {}
+        self.fail_on_role = fail_on_role
+        self.non_existent_roles = []
+        self.changed = False
+        self.__check_roles_exist()
+
+    def grant(self):
+        for group in self.groups:
+            self.granted[group] = []
+
+            for role in self.target_roles:
+                # If role is in a group now, pass:
+                if self.__check_membership(group, role):
+                    continue
+
+                query = "GRANT %s TO %s" % ((pg_quote_identifier(group, 'role'),
+                                            (pg_quote_identifier(role, 'role'))))
+                self.changed = exec_sql(self, query, ddl=True)
+
+                if self.changed:
+                    self.granted[group].append(role)
+
+        return self.changed
+
+    def revoke(self):
+        for group in self.groups:
+            self.revoked[group] = []
+
+            for role in self.target_roles:
+                # If role is not in a group now, pass:
+                if not self.__check_membership(group, role):
+                    continue
+
+                query = "REVOKE %s FROM %s" % ((pg_quote_identifier(group, 'role'),
+                                               (pg_quote_identifier(role, 'role'))))
+                self.changed = exec_sql(self, query, ddl=True)
+
+                if self.changed:
+                    self.revoked[group].append(role)
+
+        return self.changed
+
+    def __check_membership(self, src_role, dst_role):
+        query = ("SELECT ARRAY(SELECT b.rolname FROM "
+                 "pg_catalog.pg_auth_members m "
+                 "JOIN pg_catalog.pg_roles b ON (m.roleid = b.oid) "
+                 "WHERE m.member = r.oid) "
+                 "FROM pg_catalog.pg_roles r "
+                 "WHERE r.rolname = '%s'" % dst_role)
+
+        res = exec_sql(self, query, add_to_executed=False)
+        membership = []
+        if res:
+            membership = res[0][0]
+
+        if not membership:
+            return False
+
+        if src_role in membership:
+            return True
+
+        return False
+
+    def __check_roles_exist(self):
+        existent_groups = self.__roles_exist(self.groups)
+        existent_roles = self.__roles_exist(self.target_roles)
+
+        for group in self.groups:
+            if group not in existent_groups:
+                if self.fail_on_role:
+                    self.module.fail_json(msg="Role %s does not exist" % group)
+                else:
+                    self.module.warn("Role %s does not exist, pass" % group)
+                    self.non_existent_roles.append(group)
+
+        for role in self.target_roles:
+            if role not in existent_roles:
+                if self.fail_on_role:
+                    self.module.fail_json(msg="Role %s does not exist" % role)
+                else:
+                    self.module.warn("Role %s does not exist, pass" % role)
+
+                if role not in self.groups:
+                    self.non_existent_roles.append(role)
+
+                else:
+                    if self.fail_on_role:
+                        self.module.exit_json(msg="Role role '%s' is a member of role '%s'" % (role, role))
+                    else:
+                        self.module.warn("Role role '%s' is a member of role '%s', pass" % (role, role))
+
+        # Update role lists, excluding non existent roles:
+        self.groups = [g for g in self.groups if g not in self.non_existent_roles]
+
+        self.target_roles = [r for r in self.target_roles if r not in self.non_existent_roles]
+
+    def __roles_exist(self, roles):
+        tmp = ["'" + x + "'" for x in roles]
+        query = "SELECT rolname FROM pg_roles WHERE rolname IN (%s)" % ','.join(tmp)
+        return [x[0] for x in exec_sql(self, query, add_to_executed=False)]

lib/ansible/modules/database/postgresql/postgresql_membership.py

@@ -138,125 +138,19 @@ except ImportError:
     pass
 
 from ansible.module_utils.basic import AnsibleModule
-from ansible.module_utils.database import pg_quote_identifier
 from ansible.module_utils.postgres import (
     connect_to_db,
     exec_sql,
     get_conn_params,
+    PgMembership,
     postgres_common_argument_spec,
 )
 
 
-class PgMembership(object):
-    def __init__(self, module, cursor, groups, target_roles, fail_on_role):
-        self.module = module
-        self.cursor = cursor
-        self.target_roles = [r.strip() for r in target_roles]
-        self.groups = [r.strip() for r in groups]
-        self.executed_queries = []
-        self.granted = {}
-        self.revoked = {}
-        self.fail_on_role = fail_on_role
-        self.non_existent_roles = []
-        self.changed = False
-        self.__check_roles_exist()
-
-    def grant(self):
-        for group in self.groups:
-            self.granted[group] = []
-
-            for role in self.target_roles:
-                # If role is in a group now, pass:
-                if self.__check_membership(group, role):
-                    continue
-
-                query = "GRANT %s TO %s" % ((pg_quote_identifier(group, 'role'),
-                                            (pg_quote_identifier(role, 'role'))))
-                self.changed = exec_sql(self, query, ddl=True)
-
-                if self.changed:
-                    self.granted[group].append(role)
-
-        return self.changed
-
-    def revoke(self):
-        for group in self.groups:
-            self.revoked[group] = []
-
-            for role in self.target_roles:
-                # If role is not in a group now, pass:
-                if not self.__check_membership(group, role):
-                    continue
-
-                query = "REVOKE %s FROM %s" % ((pg_quote_identifier(group, 'role'),
-                                               (pg_quote_identifier(role, 'role'))))
-                self.changed = exec_sql(self, query, ddl=True)
-
-                if self.changed:
-                    self.revoked[group].append(role)
-
-        return self.changed
-
-    def __check_membership(self, src_role, dst_role):
-        query = ("SELECT ARRAY(SELECT b.rolname FROM "
-                 "pg_catalog.pg_auth_members m "
-                 "JOIN pg_catalog.pg_roles b ON (m.roleid = b.oid) "
-                 "WHERE m.member = r.oid) "
-                 "FROM pg_catalog.pg_roles r "
-                 "WHERE r.rolname = '%s'" % dst_role)
-
-        res = exec_sql(self, query, add_to_executed=False)
-        membership = []
-        if res:
-            membership = res[0][0]
-
-        if not membership:
-            return False
-
-        if src_role in membership:
-            return True
-
-        return False
-
-    def __check_roles_exist(self):
-        for group in self.groups:
-            if not self.__role_exists(group):
-                if self.fail_on_role:
-                    self.module.fail_json(msg="Role %s does not exist" % group)
-                else:
-                    self.module.warn("Role %s does not exist, pass" % group)
-                    self.non_existent_roles.append(group)
-
-        for role in self.target_roles:
-            if not self.__role_exists(role):
-                if self.fail_on_role:
-                    self.module.fail_json(msg="Role %s does not exist" % role)
-                else:
-                    self.module.warn("Role %s does not exist, pass" % role)
-
-                if role not in self.groups:
-                    self.non_existent_roles.append(role)
-
-                else:
-                    if self.fail_on_role:
-                        self.module.exit_json(msg="Role role '%s' is a member of role '%s'" % (role, role))
-                    else:
-                        self.module.warn("Role role '%s' is a member of role '%s', pass" % (role, role))
-
-        # Update role lists, excluding non existent roles:
-        self.groups = [g for g in self.groups if g not in self.non_existent_roles]
-
-        self.target_roles = [r for r in self.target_roles if r not in self.non_existent_roles]
-
-    def __role_exists(self, role):
-        return exec_sql(self, "SELECT 1 FROM pg_roles WHERE rolname = '%s'" % role, add_to_executed=False)
-
-
 # ===========================================
 # Module execution.
 #
 
-
 def main():
     argument_spec = postgres_common_argument_spec()
     argument_spec.update(

lib/ansible/modules/database/postgresql/postgresql_user.py

@@ -253,6 +253,7 @@ from ansible.module_utils.postgres import (
     connect_to_db,
     exec_sql,
     get_conn_params,
+    PgMembership,
     postgres_common_argument_spec,
 )
 from ansible.module_utils._text import to_bytes, to_native
@@ -769,96 +770,10 @@ def get_valid_flags_by_version(cursor):
     ]
 
 
-class PgMembership():
-    def __init__(self, module, cursor, target_roles, groups, fail_on_role=True):
-        self.module = module
-        self.cursor = cursor
-        self.target_roles = [r.strip() for r in target_roles]
-        self.groups = groups
-        self.granted = {}
-        self.fail_on_role = fail_on_role
-        self.non_existent_roles = []
-        self.changed = False
-        self.__check_roles_exist()
-
-    def grant(self):
-        for group in self.groups:
-            self.granted[group] = []
-
-            for role in self.target_roles:
-                # If role is in a group now, pass:
-                if self.__check_membership(group, role):
-                    continue
-
-                query = "GRANT %s TO %s" % ((pg_quote_identifier(group, 'role'),
-                                            (pg_quote_identifier(role, 'role'))))
-                self.changed = exec_sql(self, query, ddl=True, add_to_executed=False)
-                executed_queries.append(query)
-
-                if self.changed:
-                    self.granted[group].append(role)
-
-        return self.changed
-
-    def __check_membership(self, src_role, dst_role):
-        query = ("SELECT ARRAY(SELECT b.rolname FROM "
-                 "pg_catalog.pg_auth_members m "
-                 "JOIN pg_catalog.pg_roles b ON (m.roleid = b.oid) "
-                 "WHERE m.member = r.oid) "
-                 "FROM pg_catalog.pg_roles r "
-                 "WHERE r.rolname = '%s'" % dst_role)
-
-        res = exec_sql(self, query, add_to_executed=False)
-        membership = []
-        if res:
-            membership = res[0][0]
-
-        if not membership:
-            return False
-
-        if src_role in membership:
-            return True
-
-        return False
-
-    def __check_roles_exist(self):
-        for group in self.groups:
-            if not self.__role_exists(group):
-                if self.fail_on_role:
-                    self.module.fail_json(msg="Role %s does not exist" % group)
-                else:
-                    self.module.warn("Role %s does not exist, pass" % group)
-                    self.non_existent_roles.append(group)
-
-        for role in self.target_roles:
-            if not self.__role_exists(role):
-                if self.fail_on_role:
-                    self.module.fail_json(msg="Role %s does not exist" % role)
-                else:
-                    self.module.warn("Role %s does not exist, pass" % role)
-
-                if role not in self.groups:
-                    self.non_existent_roles.append(role)
-
-                else:
-                    if self.fail_on_role:
-                        self.module.exit_json(msg="Role role '%s' is a member of role '%s'" % (role, role))
-                    else:
-                        self.module.warn("Role role '%s' is a member of role '%s', pass" % (role, role))
-
-        # Update role lists, excluding non existent roles:
-        self.groups = [g for g in self.groups if g not in self.non_existent_roles]
-
-        self.target_roles = [r for r in self.target_roles if r not in self.non_existent_roles]
-
-    def __role_exists(self, role):
-        return exec_sql(self, "SELECT 1 FROM pg_roles WHERE rolname = '%s'" % role, add_to_executed=False)
-
 # ===========================================
 # Module execution.
 #
 
-
 def main():
     argument_spec = postgres_common_argument_spec()
     argument_spec.update(
@@ -938,8 +853,9 @@ def main():
         if groups:
             target_roles = []
             target_roles.append(user)
-            pg_membership = PgMembership(module, cursor, target_roles, groups)
+            pg_membership = PgMembership(module, cursor, groups, target_roles)
             changed = pg_membership.grant()
+            executed_queries.extend(pg_membership.executed_queries)
 
     else:
         if user_exists(cursor, user):