From 87340f0650e411b7e51ffa4087afddef58b9a9d6 Mon Sep 17 00:00:00 2001
From: Andrew Klychkov <aaklychkov@mail.ru>
Date: Fri, 16 Aug 2024 13:49:26 +0200
Subject: [PATCH] add roles argument

---
 plugins/modules/clickhouse_user.py            | 50 ++++++++++++++++---
 .../targets/clickhouse_user/tasks/initial.yml | 17 ++++++-
 2 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/plugins/modules/clickhouse_user.py b/plugins/modules/clickhouse_user.py
index f4183ce..02b6f75 100644
--- a/plugins/modules/clickhouse_user.py
+++ b/plugins/modules/clickhouse_user.py
@@ -82,9 +82,18 @@
     type: list
     elements: str
     version_added: '0.5.0'
+  roles:
+    description:
+      - Grants specified roles for the user.
+    type: list
+    elements: str
+    version_added: '0.6.0'
   default_roles:
     description:
-      - Grants and sets the role(s) as default for the user.
+      - Sets specified roles as default for the user.
+      - The roles must be explicitly granted to the user whether manually
+        before using this argument or by using the I(roles)
+        argument in the same task.
     type: list
     elements: str
     version_added: '0.6.0'
@@ -100,9 +109,11 @@
     name: test_user
     password: qwerty
     type_password: sha256_password
-    default_roles:
+    roles:
     - accountant
     - manager
+    default_roles:
+    - accountant
 
 - name: If user exists, update password
   community.clickhouse.clickhouse_user:
@@ -173,7 +184,10 @@ def __init__(self, module, client, name, password, type_password, cluster):
         self.cluster = cluster
         # Set default values, then update
         self.user_exists = False
-        self.default_roles_list = []
+        self.current_default_roles = []
+        self.current_roles = []
+        # Fetch actual values from DB and
+        # update the attributes with them
         self.__populate_info()
 
     def __populate_info(self):
@@ -191,7 +205,16 @@ def __populate_info(self):
 
         if result != []:
             self.user_exists = True
-            self.default_roles_list = result[0][3]
+            self.current_default_roles = result[0][3]
+
+        if self.user_exists:
+            self.current_roles = self.__fetch_user_groups()
+
+    def __fetch_user_groups(self):
+        query = ("SELECT granted_role_name FROM system.role_grants "
+                 "WHERE user_name = '%s'" % self.name)
+        result = execute_query(self.module, self.client, query)
+        return [row[0] for row in result]
 
     def create(self):
         list_settings = self.module.params['settings']
@@ -216,23 +239,35 @@ def create(self):
         if not self.module.check_mode:
             execute_query(self.module, self.client, query)
 
+        if self.module.params['roles']:
+            self.__grant_role(self.module.params['roles'])
+
         if self.module.params['default_roles']:
-            self.__grant_role(self.module.params['default_roles'])
             self.__set_default_roles(self.module.params['default_roles'])
 
         return True
 
     def update(self, update_password):
+        if self.module.params['roles']:
+            desired_roles = self.module.params['roles']
+
+            roles_to_set = []
+            for role in desired_roles:
+                if role not in self.current_roles:
+                    roles_to_set.append(role)
+
+            if roles_to_set:
+                self.__grant_roles(roles_to_set)
+
         if self.module.params['default_roles']:
             default_roles = self.module.params['default_roles']
 
             roles_to_set = []
             for role in default_roles:
-                if role not in self.default_roles_list:
+                if role not in self.current_default_roles:
                     roles_to_set.append(role)
 
             if roles_to_set:
-                self.__grant_roles(roles_to_set)
                 self.__set_default_roles(roles_to_set)
 
         if update_password == 'on_create':
@@ -296,6 +331,7 @@ def main():
             default='on_create', no_log=False
         ),
         settings=dict(type='list', elements='str'),
+        roles=dict(type='list', elements='str', default=None),
         default_roles=dict(type='list', elements='str', default=None),
     )
 
diff --git a/tests/integration/targets/clickhouse_user/tasks/initial.yml b/tests/integration/targets/clickhouse_user/tasks/initial.yml
index ba8aa1a..d322e82 100644
--- a/tests/integration/targets/clickhouse_user/tasks/initial.yml
+++ b/tests/integration/targets/clickhouse_user/tasks/initial.yml
@@ -145,6 +145,9 @@
   community.clickhouse.clickhouse_user:
     state: present
     name: test_user
+    roles:
+    - accountant
+    - manager
     default_roles:
     - accountant
     - manager
@@ -171,6 +174,9 @@
   community.clickhouse.clickhouse_user:
     state: present
     name: test_user
+    roles:
+    - accountant
+    - manager
     default_roles:
     - accountant
     - manager
@@ -198,6 +204,8 @@
   community.clickhouse.clickhouse_user:
     state: present
     name: test_user
+    roles:
+    - sales
     default_roles:
     - sales
 
@@ -223,6 +231,7 @@
   community.clickhouse.clickhouse_user:
     state: present
     name: test_user
+    roles: sales
     default_roles: sales
 
 - name: Check ret values 
@@ -247,6 +256,8 @@
   community.clickhouse.clickhouse_user:
     state: present
     name: test_user
+    roles:
+    - sales
     default_roles:
     - sales
 
@@ -262,6 +273,8 @@
   community.clickhouse.clickhouse_user:
     state: present
     name: test_user
+    roles:
+    - sales
     default_roles:
     - sales
 
@@ -276,6 +289,8 @@
   community.clickhouse.clickhouse_user:
     state: present
     name: test_user
+    roles:
+    - accountant
     default_roles:
     - accountant
 
@@ -283,7 +298,7 @@
   ansible.builtin.assert:
     that:
     - result is changed
-    - result.executed_statements == ["GRANT accountant TO test_user", "ALTER USER test_user DEFAULT ROLE accountant"]
+    - result.executed_statements == ["ALTER USER test_user DEFAULT ROLE accountant"]
 
 - name: Check the actual state
   register: result