From c23573f41180794a936ba7a59c2063937cad35c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Veli=20Kerim=20=C3=87elik?= <vkcelik@gmail.com>
Date: Fri, 13 Dec 2024 13:51:50 +0100
Subject: [PATCH] allow and handle role list

---
 src/snowflake/snowpark/session.py | 32 +++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py
index 0231ad9c870..a8e20279a9a 100644
--- a/src/snowflake/snowpark/session.py
+++ b/src/snowflake/snowpark/session.py
@@ -3602,19 +3602,43 @@ def use_role(self, role: str) -> None:
         """
         self._use_object(role, "role")
 
-    def use_secondary_roles(self, roles: Optional[Literal["all", "none"]]) -> None:
+    def use_secondary_roles(self, roles: Optional[Union[List[str], Literal['all', 'none']]]) -> None:
         """
         Specifies the active/current secondary roles for the session.
         The currently-active secondary roles set the context that determines whether
         the current user has the necessary privileges to perform SQL actions.
 
+        References: `Snowflake command USE SECONDARY ROLES <https://docs.snowflake.com/en/sql-reference/sql/use-secondary-roles.html>`_.
+
         Args:
-            roles: "all" or "none". ``None`` means "none".
+            roles: list of specific roles or "all" or "none". ``None`` means "none".
 
-        References: `Snowflake command USE SECONDARY ROLES <https://docs.snowflake.com/en/sql-reference/sql/use-secondary-roles.html>`_.
+        Example 1
+            Use specific roles as secondary roles:
+
+            >>> secondary_roles = ['READER_ROLE_A', 'READER_ROLE_B']
+            >>> session.use_secondary_roles(secondary_roles)
+
+        Example 2
+            Use all roles that have been granted in addition to the current active primary role:
+
+            >>> session.use_secondary_roles('all')
+
+        Example 3
+            Use no secondary roles. Disables secondary roles in the current session:
+
+            >>> session.use_secondary_roles('none')
         """
+        def format_roles(value):
+            if not roles:
+                return 'none'
+            if isinstance(value, List):
+                # format the list according to syntax: <role_name> [ , <role_name> ... ]
+                return ', '.join(roles)
+            return value.lower()
+
         self._run_query(
-            f"use secondary roles {'none' if roles is None else roles.lower()}"
+            f"use secondary roles {format_roles(roles)}"
         )
 
     def _use_object(self, object_name: str, object_type: str) -> None: