Snow 2112354 fix empty current object (#2330)

* fix sql-facade

* fix sql_executor

Author: sfc-gh-pczajka

RELEASE-NOTES.md

@@ -31,6 +31,7 @@
 * Added `--encryption` flag to `snow stage create` command defining the type of encryption for all files on the stage.
 
 ## Fixes and improvements
+* Fix `use` commands error if current database is not set.
 
 
 # v3.8.3

src/snowflake/cli/_plugins/nativeapp/sf_sql_facade.py

@@ -159,7 +159,10 @@ def _use_object_optional(self, object_type: UseObjectType, name: str | None):
             )
 
         try:
-            prev_obj = to_identifier(current_obj_result_row[0])
+            if current_obj_result_row[0]:
+                prev_obj = to_identifier(current_obj_result_row[0])
+            else:
+                prev_obj = None
         except IndexError:
             prev_obj = None
 

src/snowflake/cli/api/sql_execution.py

@@ -124,8 +124,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):
@@ -135,14 +137,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",
     [

tests/nativeapp/test_sf_sql_facade.py

@@ -841,6 +841,24 @@ def test_use_role_same_id(mock_execute_query, mock_cursor, new_role, current_rol
     assert mock_execute_query.mock_calls == expected
 
 
+def test_use_role_current_role_empty(mock_execute_query, mock_cursor):
+    side_effects, expected = mock_execute_helper(
+        [
+            (
+                mock_cursor([(None,)], []),
+                mock.call("select current_role()"),
+            ),
+            (None, mock.call('use role "test role"')),
+        ]
+    )
+    mock_execute_query.side_effect = side_effects
+
+    with sql_facade._use_role_optional("test role"):  # noqa: SLF001
+        pass
+
+    assert mock_execute_query.mock_calls == expected
+
+
 @pytest.mark.parametrize(
     "old_db, expected_old_db",
     [("old_db", "old_db"), ("old db", '"old db"')],
@@ -888,6 +906,24 @@ def test_use_db_same_id(mock_execute_query, mock_cursor, new_db, current_db):
     assert mock_execute_query.mock_calls == expected
 
 
+def test_use_db_current_db_empty(mock_execute_query, mock_cursor):
+    side_effects, expected = mock_execute_helper(
+        [
+            (
+                mock_cursor([(None,)], []),
+                mock.call("select current_database()"),
+            ),
+            (None, mock.call('use database "new db"')),
+        ]
+    )
+    mock_execute_query.side_effect = side_effects
+
+    with sql_facade._use_database_optional("new db"):  # noqa: SLF001
+        pass
+
+    assert mock_execute_query.mock_calls == expected
+
+
 @pytest.mark.parametrize(
     "old_schema, expected_old_schema",
     [("old_schema", "old_schema"), ("old schema", '"old schema"')],
@@ -941,6 +977,24 @@ def test_use_schema_same_id(
     assert mock_execute_query.mock_calls == expected
 
 
+def test_use_schema_current_schema_empty(mock_execute_query, mock_cursor):
+    side_effects, expected = mock_execute_helper(
+        [
+            (
+                mock_cursor([(None,)], []),
+                mock.call("select current_schema()"),
+            ),
+            (None, mock.call('use schema "new schema"')),
+        ]
+    )
+    mock_execute_query.side_effect = side_effects
+
+    with sql_facade._use_schema_optional("new schema"):  # noqa: SLF001
+        pass
+
+    assert mock_execute_query.mock_calls == expected
+
+
 @pytest.mark.parametrize(
     "error_raised, error_caught, error_message",
     [