From 78929f96343290d2cec2c903a587b94301c0df6e Mon Sep 17 00:00:00 2001
From: Andrew Klychkov <aaklychkov@mail.ru>
Date: Fri, 9 Aug 2024 10:06:08 +0200
Subject: [PATCH] Use alter user to assign multiple roles as default

---
 plugins/modules/clickhouse_user.py            | 23 +++++++++++--------
 .../targets/clickhouse_user/tasks/initial.yml | 17 ++++++++------
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/plugins/modules/clickhouse_user.py b/plugins/modules/clickhouse_user.py
index 380b443..f4183ce 100644
--- a/plugins/modules/clickhouse_user.py
+++ b/plugins/modules/clickhouse_user.py
@@ -226,10 +226,14 @@ def update(self, update_password):
         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:
-                    self.__grant_role(role)
-                    self.__set_default_roles(role)
+                    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':
             return False or self.changed
@@ -259,17 +263,18 @@ def drop(self):
 
         return True
 
-    def __grant_role(self, role):
-        query = "GRANT %s TO %s" % (role, self.name)
-        executed_statements.append(query)
+    def __grant_roles(self, roles_to_set):
+        for role in roles_to_set:
+            query = "GRANT %s TO %s" % (role, self.name)
+            executed_statements.append(query)
 
-        if not self.module.check_mode:
-            execute_query(self.module, self.client, query)
+            if not self.module.check_mode:
+                execute_query(self.module, self.client, query)
 
         self.changed = True
 
-    def __set_default_roles(self, role):
-        query = "SET DEFAULT ROLE %s TO %s" % (role, self.name)
+    def __set_default_roles(self, roles_to_set):
+        query = "ALTER USER %s DEFAULT ROLE %s" % (self.name, ', '.join(roles_to_set))
         executed_statements.append(query)
 
         if not self.module.check_mode:
diff --git a/tests/integration/targets/clickhouse_user/tasks/initial.yml b/tests/integration/targets/clickhouse_user/tasks/initial.yml
index b4d1751..ba8aa1a 100644
--- a/tests/integration/targets/clickhouse_user/tasks/initial.yml
+++ b/tests/integration/targets/clickhouse_user/tasks/initial.yml
@@ -134,6 +134,7 @@
 - name: Create test roles
   loop:
   - accountant
+  - manager
   - sales
   community.clickhouse.clickhouse_role:
     name: "{{ item }}"
@@ -146,12 +147,13 @@
     name: test_user
     default_roles:
     - accountant
+    - manager
 
 - name: Check ret values 
   ansible.builtin.assert:
     that:
     - result is changed
-    - result.executed_statements == ["GRANT accountant TO test_user", "SET DEFAULT ROLE accountant TO test_user"]
+    - result.executed_statements == ["GRANT accountant TO test_user", "GRANT manager TO test_user", "ALTER USER test_user DEFAULT ROLE accountant, manager"]
 
 - name: Check the actual state
   register: result
@@ -171,12 +173,13 @@
     name: test_user
     default_roles:
     - accountant
+    - manager
 
 - name: Check ret values 
   ansible.builtin.assert:
     that:
     - result is changed
-    - result.executed_statements == ["GRANT accountant TO test_user", "SET DEFAULT ROLE accountant TO test_user"]
+    - result.executed_statements == ["GRANT accountant TO test_user", "GRANT manager TO test_user", "ALTER USER test_user DEFAULT ROLE accountant, manager"]
 
 - name: Check the actual state
   register: result
@@ -187,7 +190,7 @@
   ansible.builtin.assert:
     that:
     - result is not changed
-    - result["users"]["test_user"]["default_roles_list"] == ["accountant"]
+    - result["users"]["test_user"]["default_roles_list"] == ["accountant", "manager"]
 
 - name: Set another role as default in check mode
   register: result
@@ -202,7 +205,7 @@
   ansible.builtin.assert:
     that:
     - result is changed
-    - result.executed_statements == ["GRANT sales TO test_user", "SET DEFAULT ROLE sales TO test_user"]
+    - result.executed_statements == ["GRANT sales TO test_user", "ALTER USER test_user DEFAULT ROLE sales"]
 
 - name: Check the actual state
   register: result
@@ -213,7 +216,7 @@
   ansible.builtin.assert:
     that:
     - result is not changed
-    - result["users"]["test_user"]["default_roles_list"] == ["accountant"]
+    - result["users"]["test_user"]["default_roles_list"] == ["accountant", "manager"]
 
 - name: Set another role as default in real mode
   register: result
@@ -226,7 +229,7 @@
   ansible.builtin.assert:
     that:
     - result is changed
-    - result.executed_statements == ["GRANT sales TO test_user", "SET DEFAULT ROLE sales TO test_user"]
+    - result.executed_statements == ["GRANT sales TO test_user", "ALTER USER test_user DEFAULT ROLE sales"]
 
 - name: Check the actual state
   register: result
@@ -280,7 +283,7 @@
   ansible.builtin.assert:
     that:
     - result is changed
-    - result.executed_statements == ["GRANT accountant TO test_user", "SET DEFAULT ROLE accountant TO test_user"]
+    - result.executed_statements == ["GRANT accountant TO test_user", "ALTER USER test_user DEFAULT ROLE accountant"]
 
 - name: Check the actual state
   register: result