fix sql_executor

Author: sfc-gh-pczajka

src/snowflake/cli/api/sql_execution.py

@@ -126,8 +126,10 @@ def use(self, object_type: ObjectType, name: str):
             # Rewrite the error to make the message more useful.
             raise CouldNotUseObjectError(object_type=object_type, name=name) from err
 
-    def current_role(self) -> str:
-        return to_identifier(self.execute_query(f"select current_role()").fetchone()[0])
+    def current_role(self) -> Optional[str]:
+        if result := self.execute_query(f"select current_role()").fetchone()[0]:
+            return to_identifier(result)
+        return None
 
     @contextmanager
     def use_role(self, new_role: str):
@@ -137,14 +139,17 @@ def use_role(self, new_role: str):
         """
         new_role = to_identifier(new_role)
         prev_role = self.current_role()
-        is_different_role = new_role.lower() != prev_role.lower()
+        if prev_role:
+            is_different_role = new_role.lower() != prev_role.lower()
+        else:
+            is_different_role = True
         if is_different_role:
             self._log.debug("Assuming different role: %s", new_role)
             self.execute_query(f"use role {new_role}")
         try:
             yield
         finally:
-            if is_different_role:
+            if is_different_role and prev_role:
                 self.execute_query(f"use role {prev_role}")
 
     def session_has_warehouse(self) -> bool:

tests/api/test_sql_execution.py

@@ -63,6 +63,17 @@ def test_use_role_same_id(mock_execute_query, mock_cursor, new_role, current_rol
     assert mock_execute_query.mock_calls == [mock.call("select current_role()")]
 
 
+@mock.patch(EXECUTE_QUERY)
+def test_use_role_no_current_role(mock_execute_query, mock_cursor):
+    mock_execute_query.return_value = mock_cursor([(None,)], [])
+    with SqlExecutor().use_role("new_role"):
+        pass
+    assert mock_execute_query.mock_calls == [
+        mock.call("select current_role()"),
+        mock.call(f"use role new_role"),
+    ]
+
+
 @pytest.mark.parametrize(
     "new_warehouse, expected_new_warehouse, current_warehouse, expected_current_warehouse",
     [
@@ -94,6 +105,20 @@ def test_use_warehouse_different_id(
     ]
 
 
+@mock.patch(EXECUTE_QUERY)
+def test_use_warehouse_no_current_wh(
+    mock_execute_query,
+    mock_cursor,
+):
+    mock_execute_query.return_value = mock_cursor([(None,)], [])
+    with SqlExecutor().use_warehouse("new_warehouse"):
+        pass
+    assert mock_execute_query.mock_calls == [
+        mock.call("select current_warehouse()"),
+        mock.call(f"use warehouse new_warehouse"),
+    ]
+
+
 @pytest.mark.parametrize(
     "new_warehouse, current_warehouse",
     [