SNOW-1011769: Convenience command 'spcs image-registry url' to get registry url based on current account.

RELEASE-NOTES.md

@@ -5,6 +5,7 @@
 ## New additions
 * Added ability to specify scope of the `object list` command with the `--in <scope_type> <scope_name>` option.
 * Introduced `snowflake.cli.api.console.cli_console` object with helper methods for intermediate output.
+* Added new convenience command `spcs image-registry url` to get the URL for your account image registry.
 * Added convenience function `spcs image-repository url <repo_name>`.
 
 ## Fixes and improvements

src/snowflake/cli/plugins/spcs/image_registry/commands.py

@@ -1,11 +1,15 @@
 import typer
+from click import ClickException
 from snowflake.cli.api.commands.decorators import (
     global_options_with_connection,
     with_output,
 )
 from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS
-from snowflake.cli.api.output.types import ObjectResult
-from snowflake.cli.plugins.spcs.image_registry.manager import RegistryManager
+from snowflake.cli.api.output.types import MessageResult, ObjectResult
+from snowflake.cli.plugins.spcs.image_registry.manager import (
+    NoImageRepositoriesFoundError,
+    RegistryManager,
+)
 
 app = typer.Typer(
     context_settings=DEFAULT_CONTEXT_SETTINGS,
@@ -19,5 +23,19 @@
 @with_output
 @global_options_with_connection
 def token(**options) -> ObjectResult:
-    """Gets the token from environment to use for authenticating with the registry."""
+    """Gets the token from environment to use for authenticating with the registry. Note that this token is specific
+    to your current user and will not grant access to any repositories that your current user cannot access."""
     return ObjectResult(RegistryManager().get_token())
+
+
+@app.command()
+@with_output
+@global_options_with_connection
+def url(**options) -> MessageResult:
+    """Gets the image registry URL for the current account. Must be called from a role that can view at least one image repository in the image registry."""
+    try:
+        return MessageResult(RegistryManager().get_registry_url())
+    except NoImageRepositoriesFoundError:
+        raise ClickException(
+            "No image repository found. To get the registry url, please switch to a role with read access to at least one image repository or create a new image repository first."
+        )

src/snowflake/cli/plugins/spcs/image_registry/manager.py

@@ -1,10 +1,17 @@
 import base64
 import json
+import re
 from urllib.parse import urlparse
 
 import requests
 from click import ClickException
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
+from snowflake.connector.cursor import DictCursor
+
+
+class NoImageRepositoriesFoundError(ClickException):
+    def __init__(self, msg: str = "No image repository found."):
+        super().__init__(msg)
 
 
 class RegistryManager(SqlExecutionMixin):
@@ -45,3 +52,17 @@ def login_to_registry(self, repo_url):
         if resp.status_code != 200:
             raise ClickException(f"Failed to login to the repository {resp.text}")
         return json.loads(resp.text)["token"]
+
+    def _has_url_scheme(self, url: str):
+        return re.fullmatch(r"^.*//.+", url) is not None
+
+    def get_registry_url(self):
+        repositories_query = "show image repositories in account"
+        result_set = self._execute_query(repositories_query, cursor_class=DictCursor)
+        results = result_set.fetchall()
+        if len(results) == 0:
+            raise NoImageRepositoriesFoundError()
+        sample_repository_url = results[0]["repository_url"]
+        if not self._has_url_scheme(sample_repository_url):
+            sample_repository_url = f"//{sample_repository_url}"
+        return urlparse(sample_repository_url).netloc

tests/spcs/__snapshots__/test_registry.ambr

@@ -0,0 +1,11 @@
+# serializer version: 1
+# name: test_get_registry_url_no_repositories_cli
+  '''
+  ╭─ Error ──────────────────────────────────────────────────────────────────────╮
+  │ No image repository found. To get the registry url, please switch to a role  │
+  │ with read access to at least one image repository or create a new image      │
+  │ repository first.                                                            │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+
+  '''
+# ---

tests/spcs/test_registry.py

@@ -1,13 +1,17 @@
 import json
-
 from tests.testing_utils.fixtures import *
+from snowflake.cli.plugins.spcs.image_registry.manager import (
+    RegistryManager,
+    NoImageRepositoriesFoundError,
+)
+from snowflake.connector.cursor import DictCursor
 
 
 @mock.patch("snowflake.cli.plugins.spcs.image_registry.manager.RegistryManager._conn")
 @mock.patch(
     "snowflake.cli.plugins.spcs.image_registry.manager.RegistryManager._execute_query"
 )
-def test_registry_get_token_2(mock_execute, mock_conn, mock_cursor, runner):
+def test_registry_get_token(mock_execute, mock_conn, mock_cursor, runner):
     mock_execute.return_value = mock_cursor(
         ["row"], ["Statement executed successfully"]
     )
@@ -20,3 +24,84 @@ def test_registry_get_token_2(mock_execute, mock_conn, mock_cursor, runner):
     result = runner.invoke(["spcs", "image-registry", "token", "--format", "JSON"])
     assert result.exit_code == 0, result.output
     assert json.loads(result.stdout) == {"token": "token1234", "expires_in": 42}
+
+
+MOCK_REPO_COLUMNS = [
+    "created_on",
+    "name",
+    "database_name",
+    "schema_name",
+    "repository_url",
+    "owner",
+    "owner_role_type",
+    "comment",
+]
+
+
+@mock.patch("snowflake.cli.plugins.spcs.image_registry.manager.RegistryManager._conn")
+@mock.patch(
+    "snowflake.cli.plugins.spcs.image_registry.manager.RegistryManager._execute_query"
+)
+def test_get_registry_url(mock_execute, mock_conn, mock_cursor):
+    mock_row = [
+        "2023-01-01 00:00:00",
+        "IMAGES",
+        "DB",
+        "SCHEMA",
+        "orgname-alias.registry.snowflakecomputing.com/DB/SCHEMA/IMAGES",
+        "TEST_ROLE",
+        "ROLE",
+        "",
+    ]
+
+    mock_execute.return_value = mock_cursor(
+        rows=[{col: row for col, row in zip(MOCK_REPO_COLUMNS, mock_row)}],
+        columns=MOCK_REPO_COLUMNS,
+    )
+    result = RegistryManager().get_registry_url()
+    expected_query = "show image repositories in account"
+    mock_execute.assert_called_once_with(expected_query, cursor_class=DictCursor)
+    assert result == "orgname-alias.registry.snowflakecomputing.com"
+
+
+@mock.patch("snowflake.cli.plugins.spcs.image_registry.manager.RegistryManager._conn")
+@mock.patch(
+    "snowflake.cli.plugins.spcs.image_registry.manager.RegistryManager._execute_query"
+)
+def test_get_registry_url_no_repositories(mock_execute, mock_conn, mock_cursor):
+    mock_execute.return_value = mock_cursor(
+        rows=[],
+        columns=MOCK_REPO_COLUMNS,
+    )
+    with pytest.raises(NoImageRepositoriesFoundError):
+        RegistryManager().get_registry_url()
+
+    expected_query = "show image repositories in account"
+    mock_execute.assert_called_once_with(expected_query, cursor_class=DictCursor)
+
+
+@mock.patch(
+    "snowflake.cli.plugins.spcs.image_registry.manager.RegistryManager.get_registry_url"
+)
+def test_get_registry_url_no_repositories_cli(mock_get_registry_url, runner, snapshot):
+    mock_get_registry_url.side_effect = NoImageRepositoriesFoundError()
+    result = runner.invoke(["spcs", "image-registry", "url"])
+    assert result.exit_code == 1, result.output
+    assert result.output == snapshot
+
+
+@pytest.mark.parametrize(
+    "url, expected",
+    [
+        ("www.google.com", False),
+        ("https://www.google.com", True),
+        ("//www.google.com", True),
+        ("snowservices.registry.snowflakecomputing.com/db/schema/tutorial_repo", False),
+        (
+            "http://snowservices.registry.snowflakecomputing.com/db/schema/tutorial_repo",
+            True,
+        ),
+    ],
+)
+def test_has_url_scheme(url: str, expected: bool):
+    assert RegistryManager()._has_url_scheme(url) == expected

tests_integration/snowflake_connector.py

@@ -24,6 +24,16 @@ def test_database(snowflake_session):
     del os.environ[f"{_ENV_PARAMETER_PREFIX}_DATABASE"]
 
 
+@pytest.fixture(scope="function")
+def test_role(snowflake_session):
+    role_name = f"role_{uuid.uuid4().hex}"
+    snowflake_session.execute_string(
+        f"create role {role_name}; grant role {role_name} to user {snowflake_session.user};"
+    )
+    yield role_name
+    snowflake_session.execute_string(f"drop role {role_name}")
+
+
 @pytest.fixture(scope="session")
 def snowflake_session():
     config = {

tests_integration/spcs/test_registry.py

@@ -1,5 +1,8 @@
 import pytest
 
+from tests_integration.test_utils import row_from_snowflake_session
+from tests_integration.testing_utils.naming_utils import ObjectNameProvider
+
 
 @pytest.mark.integration
 def test_token(runner):
@@ -11,3 +14,33 @@ def test_token(runner):
     assert result.json["token"]
     assert "expires_in" in result.json
     assert result.json["expires_in"]
+
+
+@pytest.mark.integration
+def test_get_registry_url(test_database, test_role, runner, snowflake_session):
+    # newly created test_role should have no access to image repositories and should not be able to get registry URL
+    test_repo = ObjectNameProvider("test_repo").create_and_get_next_object_name()
+    snowflake_session.execute_string(f"create image repository {test_repo}")
+
+    fail_result = runner.invoke_with_connection(
+        ["spcs", "image-registry", "url", "--role", test_role]
+    )
+    assert fail_result.exit_code == 1, fail_result.output
+    assert "No image repository found." in fail_result.output
+
+    # role should be able to get registry URL once granted read access to an image repository
+    repo_list_cursor = snowflake_session.execute_string("show image repositories")
+    expected_repo_url = row_from_snowflake_session(repo_list_cursor)[0][
+        "repository_url"
+    ]
+    expected_registry_url = "/".join(expected_repo_url.split("/")[:-3])
+    snowflake_session.execute_string(
+        f"grant usage on database {snowflake_session.database} to role {test_role};"
+        f"grant usage on schema {snowflake_session.schema} to role {test_role};"
+        f"grant read on image repository {test_repo} to role {test_role};"
+    )
+    success_result = runner.invoke_with_connection(
+        ["spcs", "image-registry", "url", "--role", test_role]
+    )
+    assert success_result.exit_code == 0, success_result.output
+    assert success_result.output.strip() == expected_registry_url