diff --git a/src/snowflake/cli/_plugins/connection/util.py b/src/snowflake/cli/_plugins/connection/util.py index 45c6873e0a..be51927566 100644 --- a/src/snowflake/cli/_plugins/connection/util.py +++ b/src/snowflake/cli/_plugins/connection/util.py @@ -112,14 +112,16 @@ def is_regionless_redirect(conn: SnowflakeConnection) -> bool: def get_host_region(host: str) -> str | None: """ Looks for hosts of form - .[x.y.z].snowflakecomputing.com - Returns the three-part [region identifier] or None. + .[x.y[.z]].snowflakecomputing.com + Returns the two-part or three-part [region identifier] or None. """ host_parts = host.split(".") if host_parts[-1] == "local": return LOCAL_DEPLOYMENT_REGION elif len(host_parts) == 6: return ".".join(host_parts[1:4]) + elif len(host_parts) == 5: + return ".".join(host_parts[1:3]) return None @@ -172,7 +174,17 @@ def get_context(conn: SnowflakeConnection) -> str: return get_region(conn) -def get_account(conn: SnowflakeConnection) -> str: +def get_account_locator(conn: SnowflakeConnection) -> str: + """ + Determines the account locator that this connection refers to. + """ + *_, cursor = conn.execute_string( + f"select current_account_locator()", cursor_class=DictCursor + ) + return cursor.fetchone()["CURRENT_ACCOUNT_LOCATOR()"].lower() + + +def get_account_name(conn: SnowflakeConnection) -> str: """ Determines the account that this connection refers to. """ @@ -214,9 +226,12 @@ def make_snowsight_url(conn: SnowflakeConnection, path: str) -> str: """ snowsight_host = get_snowsight_host(conn) deployment = get_context(conn) - account = get_account(conn) path_with_slash = path if path.startswith("/") else f"/{path}" - return f"{snowsight_host}/{deployment}/{account}{path_with_slash}" + if len(deployment.split(".")) == 2: + account_locator = get_account_locator(conn) + return f"{snowsight_host}/{deployment}/{account_locator}{path_with_slash}" + account_name = get_account_name(conn) + return f"{snowsight_host}/{deployment}/{account_name}{path_with_slash}" def strip_if_value_present(value: Optional[str]) -> Optional[str]: diff --git a/tests/nativeapp/test_manager.py b/tests/nativeapp/test_manager.py index 91f26051d0..57348c5fe3 100644 --- a/tests/nativeapp/test_manager.py +++ b/tests/nativeapp/test_manager.py @@ -720,7 +720,7 @@ def test_get_existing_app_pkg_info_app_pkg_does_not_exist( # With connection warehouse, with PDF warehouse # Without connection warehouse, with PDF warehouse @mock.patch("snowflake.cli._plugins.connection.util.get_context") -@mock.patch("snowflake.cli._plugins.connection.util.get_account") +@mock.patch("snowflake.cli._plugins.connection.util.get_account_name") @mock.patch("snowflake.cli._plugins.connection.util.get_snowsight_host") @mock.patch(SQL_EXECUTOR_EXECUTE) @mock_connection() @@ -788,7 +788,7 @@ def test_get_snowsight_url_with_pdf_warehouse( # With connection warehouse, without PDF warehouse # Without connection warehouse, without PDF warehouse @mock.patch("snowflake.cli._plugins.connection.util.get_context") -@mock.patch("snowflake.cli._plugins.connection.util.get_account") +@mock.patch("snowflake.cli._plugins.connection.util.get_account_name") @mock.patch("snowflake.cli._plugins.connection.util.get_snowsight_host") @mock.patch(SQL_EXECUTOR_EXECUTE) @mock_connection() diff --git a/tests/streamlit/streamlit_test_class.py b/tests/streamlit/streamlit_test_class.py index b6f3ab929a..ccd5f4f932 100644 --- a/tests/streamlit/streamlit_test_class.py +++ b/tests/streamlit/streamlit_test_class.py @@ -39,7 +39,7 @@ def setup_method(self): ).start() self.mock_get_account = mock.patch( - "snowflake.cli._plugins.connection.util.get_account" + "snowflake.cli._plugins.connection.util.get_account_name" ).start() self.mock_get_account.return_value = "my_account" diff --git a/tests/test_utils.py b/tests/test_utils.py index ab1748cb99..950bfc6455 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -176,7 +176,7 @@ def test_identifier_for_url(identifier, expected): assert identifier_for_url(identifier) == expected -@patch("snowflake.cli._plugins.connection.util.get_account") +@patch("snowflake.cli._plugins.connection.util.get_account_name") @patch("snowflake.cli._plugins.connection.util.get_context") @patch("snowflake.cli._plugins.connection.util.get_snowsight_host") @pytest.mark.parametrize( @@ -218,6 +218,27 @@ def test_make_snowsight_url( assert actual == expected +@patch( + "snowflake.cli._plugins.connection.util.is_regionless_redirect", return_value=False +) +@patch("snowflake.cli._plugins.connection.util.get_region", return_value="x.y") +@patch( + "snowflake.cli._plugins.connection.util.get_account_locator", + return_value="account_locator", +) +@patch( + "snowflake.cli._plugins.connection.util.get_snowsight_host", + return_value="https://test.snowsight.host", +) +def test_get_snowsight_host_with_account_locator( + get_snowsight_host, get_locator, get_region, is_regionless +): + assert ( + make_snowsight_url(None, "an/url/path") + == "https://test.snowsight.host/x.y/account_locator/an/url/path" + ) + + @pytest.mark.parametrize( "allowlist, expected", [ @@ -286,8 +307,8 @@ def test_get_context_local_non_regionless_gets_local_region( ("org-acct.mydns.snowflakecomputing.com", None), ("account.x.us-west-2.aws.snowflakecomputing.com", "x.us-west-2.aws"), ("naf_test_pc.us-west-2.snowflakecomputing.com", None), - ("test_account.az.int.snowflakecomputing.com", None), - ("frozenweb.prod3.external-zone.snowflakecomputing.com", None), + ("test_account.az.int.snowflakecomputing.com", "az.int"), + ("frozenweb.prod3.external-zone.snowflakecomputing.com", "prod3.external-zone"), ], ) def test_get_host_region(host, expected):