Fix account locator

Author: sfc-gh-pczajka

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
-    <account>.[x.y.z].snowflakecomputing.com
-    Returns the three-part [region identifier] or None.
+    <account>.[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]:

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):