SNOW-1011769: Adding 'spcs image-registry url' command.

RELEASE-NOTES.md

@@ -30,7 +30,7 @@
   * `pool`, `job`, `service`, and `image-registry` commands were moved from `snowpark` group to a new `spcs` group (`registry` was renamed to `image-registry`).
   * `snow spcs pool create` and `snow spcs service create` have been updated with new options to match SQL interface
   * Added new `image-repository` command group under `spcs`. Moved `list-images` and `list-tags` from `registry` to `image-repository`.
-
+  * Added new convenience command to get the URL for your account image registry `spcs image-registry url`.
 * Streamlit changes
   * `snow streamlit deploy` is requiring `snowflake.yml` project file with a Streamlit definition.
   * `snow streamlit describe` is now `snow object describe streamlit`

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 (
+    NoRepositoriesViewableError,
+    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 NoRepositoriesViewableError:
+        raise ClickException(
+            "Current role cannot view any image repositories. Please ensure that at least one repository exists in your image registry and use a role with read access to at least one image repository before retrieving registry URL."
+        )

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

@@ -5,6 +5,14 @@
 import requests
 from click import ClickException
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
+from snowflake.connector.cursor import DictCursor
+
+
+class NoRepositoriesViewableError(ClickException):
+    def __init__(
+        self, msg: str = "Current role does not have view access to any repositories"
+    ):
+        super().__init__(msg)
 
 
 class RegistryManager(SqlExecutionMixin):
@@ -45,3 +53,13 @@ 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 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 NoRepositoriesViewableError()
+        sample_repository_url = results[0]["repository_url"]
+
+        return "/".join(sample_repository_url.split("/")[:-3])

tests/spcs/__snapshots__/test_registry.ambr

@@ -0,0 +1,11 @@
+# serializer version: 1
+# name: test_get_registry_url_no_repositories_cli
+  '''
+  ╭─ Error ──────────────────────────────────────────────────────────────────────╮
+  │ Current role cannot view any image repositories. Please ensure that at least │
+  │ one repository exists in your image registry and use a role with read access │
+  │ to at least one image repository before retrieving registry URL.             │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+
+  '''
+# ---

tests/spcs/test_registry.py

@@ -1,6 +1,10 @@
 import json
-
 from tests.testing_utils.fixtures import *
+from snowflake.cli.plugins.spcs.image_registry.manager import (
+    RegistryManager,
+    NoRepositoriesViewableError,
+)
+from snowflake.connector.cursor import DictCursor
 
 
 @mock.patch("snowflake.cli.plugins.spcs.image_registry.manager.RegistryManager._conn")
@@ -20,3 +24,73 @@ 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.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_rows = [
+        "2023-01-01 00:00:00",
+        "IMAGES",
+        "DB",
+        "SCHEMA",
+        "orgname-alias.registry.snowflakecomputing.com/DB/SCHEMA/IMAGES",
+        "TEST_ROLE",
+        "ROLE",
+        "",
+    ]
+    mock_columns = [
+        "created_on",
+        "name",
+        "database_name",
+        "schema_name",
+        "repository_url",
+        "owner",
+        "owner_role_type",
+        "comment",
+    ]
+    mock_execute.return_value = mock_cursor(
+        rows=[{col: row for col, row in zip(mock_columns, mock_rows)}],
+        columns=mock_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=[
+            "created_on",
+            "name",
+            "database_name",
+            "schema_name",
+            "repository_url",
+            "owner",
+            "owner_role_type",
+            "comment",
+        ],
+    )
+    with pytest.raises(NoRepositoriesViewableError):
+        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 = NoRepositoriesViewableError()
+    result = runner.invoke(["spcs", "image-registry", "url"])
+    assert result.exit_code == 1, result.output
+    assert result.output == snapshot

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/__snapshots__/test_registry.ambr

@@ -0,0 +1,11 @@
+# serializer version: 1
+# name: test_get_registry_url
+  '''
+  ╭─ Error ──────────────────────────────────────────────────────────────────────╮
+  │ Current role cannot view any image repositories. Please ensure that at least │
+  │ one repository exists in your image registry and use a role with read access │
+  │ to at least one image repository before retrieving registry URL.             │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+
+  '''
+# ---

tests_integration/spcs/test_registry.py

@@ -1,4 +1,7 @@
 import pytest
+from snowflake.connector.cursor import DictCursor
+
+from tests_integration.testing_utils.naming_utils import ObjectNameProvider
 
 
 @pytest.mark.integration
@@ -11,3 +14,35 @@ 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, snapshot
+):
+    # newly created 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 fail_result.output == snapshot
+
+    # 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", cursor_class=DictCursor
+    )
+    expected_repo_url = repo_list_cursor[0].fetchone()["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