Remove DictCursor from use_warehouse (#1588)

Followup to #1584 for `use_warehouse`
snowflakedb · Sep 18, 2024 · 55ffab6 · 55ffab6
1 parent b87e39c
commit 55ffab6
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 69 deletions.
diff --git a/src/snowflake/cli/api/sql_execution.py b/src/snowflake/cli/api/sql_execution.py
@@ -119,9 +119,9 @@ def use_role(self, new_role: str):
 
     def session_has_warehouse(self) -> bool:
         result = self._execute_query(
-            "select current_warehouse() is not null as result", cursor_class=DictCursor
+            "select current_warehouse() is not null"
         ).fetchone()
-        return bool(result.get("RESULT"))
+        return bool(result[0])
 
     @contextmanager
     def use_warehouse(self, new_wh: str):
@@ -131,12 +131,10 @@ def use_warehouse(self, new_wh: str):
         If there is no default warehouse in the account, it will throw an error.
         """
 
-        wh_result = self._execute_query(
-            f"select current_warehouse()", cursor_class=DictCursor
-        ).fetchone()
+        wh_result = self._execute_query(f"select current_warehouse()").fetchone()
         # If user has an assigned default warehouse, prev_wh will contain a value even if the warehouse is suspended.
         try:
-            prev_wh = wh_result["CURRENT_WAREHOUSE()"]
+            prev_wh = wh_result[0]
         except:
             prev_wh = None
 

diff --git a/tests/nativeapp/test_manager.py b/tests/nativeapp/test_manager.py
@@ -640,8 +640,8 @@ def test_get_snowsight_url_with_pdf_warehouse(
     side_effects, expected = mock_execute_helper(
         [
             (
-                mock_cursor([{"CURRENT_WAREHOUSE()": warehouse}], []),
-                mock.call("select current_warehouse()", cursor_class=DictCursor),
+                mock_cursor([(warehouse,)], []),
+                mock.call("select current_warehouse()"),
             ),
             (None, mock.call("use warehouse app_warehouse")),
         ]
@@ -669,7 +669,7 @@ def test_get_snowsight_url_with_pdf_warehouse(
         (
             "napp_project_1",
             "MockWarehouse",
-            [mock.call("select current_warehouse()", cursor_class=DictCursor)],
+            [mock.call("select current_warehouse()")],
             [None],
         ),
         (
@@ -701,7 +701,7 @@ def test_get_snowsight_url_without_pdf_warehouse(
     working_dir: Path = project_definition_files[0].parent
 
     mock_execute_query.side_effect = [
-        mock_cursor([{"CURRENT_WAREHOUSE()": warehouse}], [])
+        mock_cursor([(warehouse,)], [])
     ] + fallback_side_effect
 
     native_app_manager = _get_na_manager(str(working_dir))

diff --git a/tests/nativeapp/test_package_scripts.py b/tests/nativeapp/test_package_scripts.py
@@ -29,7 +29,6 @@
 )
 from snowflake.cli.api.project.definition_manager import DefinitionManager
 from snowflake.connector import ProgrammingError
-from snowflake.connector.cursor import DictCursor
 
 from tests.nativeapp.patch_utils import mock_connection
 from tests.nativeapp.utils import (
@@ -58,13 +57,13 @@ def _get_na_manager(working_dir):
         (
             "napp_project_1",  # With connection warehouse, without PDF warehouse
             [
-                mock.call("select current_warehouse()", cursor_class=DictCursor),
+                mock.call("select current_warehouse()"),
             ],
         ),
         (
             "napp_project_with_pkg_warehouse",  # With connection warehouse, with PDF warehouse
             [
-                mock.call("select current_warehouse()", cursor_class=DictCursor),
+                mock.call("select current_warehouse()"),
                 mock.call("use warehouse myapp_pkg_warehouse"),
                 mock.call("use warehouse MockWarehouse"),
             ],
@@ -83,9 +82,7 @@ def test_package_scripts_with_conn_info(
     mock_conn.return_value = MockConnectionCtx()
     working_dir: Path = project_definition_files[0].parent
     # Only consequential for "select current_warehouse()"
-    mock_execute_query.return_value = mock_cursor(
-        [{"CURRENT_WAREHOUSE()": "MockWarehouse"}], []
-    )
+    mock_execute_query.return_value = mock_cursor([("MockWarehouse",)], [])
     native_app_manager = _get_na_manager(str(working_dir))
     native_app_manager._apply_package_scripts()  # noqa: SLF001
     assert mock_execute_query.mock_calls == expected_calls
@@ -132,7 +129,7 @@ def test_package_scripts_without_conn_info_throws_error(
 ):
     mock_conn.return_value = MockConnectionCtx(warehouse=None)
     working_dir: Path = project_definition_files[0].parent
-    mock_execute_query.return_value = mock_cursor([{"CURRENT_WAREHOUSE()": None}], [])
+    mock_execute_query.return_value = mock_cursor([(None,)], [])
     native_app_manager = _get_na_manager(str(working_dir))
     with pytest.raises(ClickException) as err:
         native_app_manager._apply_package_scripts()  # noqa: SLF001
@@ -158,12 +155,12 @@ def test_package_scripts_without_conn_info_succeeds(
 ):
     mock_conn.return_value = MockConnectionCtx(warehouse=None)
     working_dir: Path = project_definition_files[0].parent
-    mock_execute_query.return_value = mock_cursor([{"CURRENT_WAREHOUSE()": None}], [])
+    mock_execute_query.return_value = mock_cursor([(None,)], [])
     native_app_manager = _get_na_manager(str(working_dir))
     native_app_manager._apply_package_scripts()  # noqa: SLF001
 
     assert mock_execute_query.mock_calls == [
-        mock.call("select current_warehouse()", cursor_class=DictCursor),
+        mock.call("select current_warehouse()"),
         mock.call("use warehouse myapp_pkg_warehouse"),
     ]
     assert mock_execute_queries.mock_calls == [
@@ -258,7 +255,12 @@ def test_package_scripts_w_missing_warehouse_exception(
 ):
     mock_conn.return_value = MockConnectionCtx()
     mock_execute_query.side_effect = [
-        mock_cursor([{"CURRENT_WAREHOUSE()": "old_wh"}], []),
+        mock_cursor(
+            [
+                ("old_wh"),
+            ],
+            [],
+        ),
         None,
         None,
     ]
@@ -287,7 +289,12 @@ def test_package_scripts_w_warehouse_access_exception(
     mock_cursor,
 ):
     side_effects = [
-        mock_cursor([{"CURRENT_WAREHOUSE()": "old_wh"}], []),
+        mock_cursor(
+            [
+                ("old_wh"),
+            ],
+            [],
+        ),
         ProgrammingError(
             msg="Object does not exist, or operation cannot be performed.",
             errno=DOES_NOT_EXIST_OR_CANNOT_BE_PERFORMED,