Skip to content

Commit

Permalink
Remove use of DictCursor for use_role (#1584)
Browse files Browse the repository at this point in the history 
Query Redaction (an enterprise feature) also redacts column names, which breaks `DictCursor`. Let's remove it from `use_role` and use a column index since the query has a well-defined column order in its output.
  • Loading branch information
@sfc-gh-fcampbell
sfc-gh-fcampbell authored Sep 18, 2024
1 parent d77ddb6 commit 27b0a87
Show file tree
Hide file tree
Showing 9 changed files with 170 additions and 180 deletions.
11 changes: 2 additions & 9 deletions src/snowflake/cli/api/sql_execution.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,22 +98,15 @@ def use(self, object_type: ObjectType, name: str):
)

def current_role(self) -> str:
*_, cursor = self._execute_string(
"select current_role()", cursor_class=DictCursor
)
role_result = cursor.fetchone()
return role_result["CURRENT_ROLE()"]
return self._execute_query(f"select current_role()").fetchone()[0]

@contextmanager
def use_role(self, new_role: str):
"""
Switches to a different role for a while, then switches back.
This is a no-op if the requested role is already active.
"""
role_result = self._execute_query(
f"select current_role()", cursor_class=DictCursor
).fetchone()
prev_role = role_result["CURRENT_ROLE()"]
prev_role = self.current_role()
is_different_role = new_role.lower() != prev_role.lower()
if is_different_role:
self._log.debug("Assuming different role: %s", new_role)
Expand Down
48 changes: 24 additions & 24 deletions tests/nativeapp/test_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ def test_sync_deploy_root_with_stage(
temp_dir,
mock_cursor,
):
mock_execute.return_value = mock_cursor([{"CURRENT_ROLE()": "old_role"}], [])
mock_execute.return_value = mock_cursor([("old_role",)], [])
mock_diff_result = DiffResult(different=[StagePath("setup.sql")])
mock_compute_stage_diff.return_value = mock_diff_result
mock_local_diff_with_stage.return_value = None
Expand All @@ -132,7 +132,7 @@ def test_sync_deploy_root_with_stage(
)

expected = [
mock.call("select current_role()", cursor_class=DictCursor),
mock.call("select current_role()"),
mock.call("use role new_role"),
mock.call(f"create schema if not exists app_pkg.app_src"),
mock.call(
Expand Down Expand Up @@ -218,8 +218,8 @@ def test_get_app_pkg_distribution_in_snowflake(mock_execute, temp_dir, mock_curs
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(
Expand Down Expand Up @@ -259,8 +259,8 @@ def test_get_app_pkg_distribution_in_snowflake_throws_programming_error(
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(
Expand Down Expand Up @@ -297,8 +297,8 @@ def test_get_app_pkg_distribution_in_snowflake_throws_execution_error(
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(mock_cursor([], []), mock.call("describe application package app_pkg")),
Expand Down Expand Up @@ -329,8 +329,8 @@ def test_get_app_pkg_distribution_in_snowflake_throws_distribution_error(
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(
Expand Down Expand Up @@ -442,8 +442,8 @@ def test_get_existing_app_info_app_exists(mock_execute, temp_dir, mock_cursor):
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role app_role")),
(
Expand Down Expand Up @@ -484,8 +484,8 @@ def test_get_existing_app_info_app_does_not_exist(mock_execute, temp_dir, mock_c
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role app_role")),
(
Expand Down Expand Up @@ -515,8 +515,8 @@ def test_get_existing_app_pkg_info_app_pkg_exists(mock_execute, temp_dir, mock_c
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(
Expand Down Expand Up @@ -562,8 +562,8 @@ def test_get_existing_app_pkg_info_app_pkg_does_not_exist(
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(
Expand Down Expand Up @@ -741,8 +741,8 @@ def test_create_app_pkg_no_existing_package(
side_effects, expected = mock_execute_helper(
[
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(
Expand Down Expand Up @@ -1162,8 +1162,8 @@ def test_validate_use_scratch_stage(
),
),
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(
Expand Down Expand Up @@ -1224,8 +1224,8 @@ def test_validate_failing_drops_scratch_stage(
),
),
(
mock_cursor([{"CURRENT_ROLE()": "old_role"}], []),
mock.call("select current_role()", cursor_class=DictCursor),
mock_cursor([("old_role",)], []),
mock.call("select current_role()"),
),
(None, mock.call("use role package_role")),
(
Expand Down
4 changes: 1 addition & 3 deletions tests/nativeapp/test_project_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -242,9 +242,7 @@ def test_project_model_explicit_package_app_name_with_suffix(
def test_project_model_falls_back_to_current_role(
mock_connect, project_definition_files: List[Path], mock_ctx, mock_cursor
):
ctx = mock_ctx(
cursor=mock_cursor([{"CURRENT_ROLE()": CURRENT_ROLE}], []), role=None
)
ctx = mock_ctx(cursor=mock_cursor([(CURRENT_ROLE,)], []), role=None)
mock_connect.return_value = ctx

project_defn = load_project(project_definition_files).project_definition
Expand Down
Loading

0 comments on commit 27b0a87

Please sign in to comment.