diff --git a/src/snowflake/cli/plugins/spcs/image_registry/commands.py b/src/snowflake/cli/plugins/spcs/image_registry/commands.py index 32e69b8614..8468a046c2 100644 --- a/src/snowflake/cli/plugins/spcs/image_registry/commands.py +++ b/src/snowflake/cli/plugins/spcs/image_registry/commands.py @@ -7,7 +7,7 @@ from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS from snowflake.cli.api.output.types import MessageResult, ObjectResult from snowflake.cli.plugins.spcs.image_registry.manager import ( - NoRepositoriesViewableError, + NoImageRepositoriesFoundError, RegistryManager, ) @@ -35,7 +35,7 @@ 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: + 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." ) diff --git a/src/snowflake/cli/plugins/spcs/image_registry/manager.py b/src/snowflake/cli/plugins/spcs/image_registry/manager.py index 55fe65b278..078bbd0e46 100644 --- a/src/snowflake/cli/plugins/spcs/image_registry/manager.py +++ b/src/snowflake/cli/plugins/spcs/image_registry/manager.py @@ -1,5 +1,6 @@ import base64 import json +import re from urllib.parse import urlparse import requests @@ -8,7 +9,7 @@ from snowflake.connector.cursor import DictCursor -class NoRepositoriesViewableError(ClickException): +class NoImageRepositoriesFoundError(ClickException): def __init__(self, msg: str = "No image repository found."): super().__init__(msg) @@ -52,12 +53,16 @@ def login_to_registry(self, repo_url): 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 NoRepositoriesViewableError() + raise NoImageRepositoriesFoundError() sample_repository_url = results[0]["repository_url"] - - return "/".join(sample_repository_url.split("/")[:-3]) + if not self._has_url_scheme(sample_repository_url): + sample_repository_url = f"//{sample_repository_url}" + return urlparse(sample_repository_url).netloc diff --git a/tests/spcs/test_registry.py b/tests/spcs/test_registry.py index c285141076..264567f174 100644 --- a/tests/spcs/test_registry.py +++ b/tests/spcs/test_registry.py @@ -2,7 +2,7 @@ from tests.testing_utils.fixtures import * from snowflake.cli.plugins.spcs.image_registry.manager import ( RegistryManager, - NoRepositoriesViewableError, + NoImageRepositoriesFoundError, ) from snowflake.connector.cursor import DictCursor @@ -11,7 +11,7 @@ @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"] ) @@ -73,7 +73,7 @@ def test_get_registry_url_no_repositories(mock_execute, mock_conn, mock_cursor): rows=[], columns=MOCK_REPO_COLUMNS, ) - with pytest.raises(NoRepositoriesViewableError): + with pytest.raises(NoImageRepositoriesFoundError): RegistryManager().get_registry_url() expected_query = "show image repositories in account" @@ -84,7 +84,24 @@ def test_get_registry_url_no_repositories(mock_execute, mock_conn, mock_cursor): "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() + 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 diff --git a/tests_integration/spcs/__snapshots__/test_registry.ambr b/tests_integration/spcs/__snapshots__/test_registry.ambr deleted file mode 100644 index a5ddc2aeed..0000000000 --- a/tests_integration/spcs/__snapshots__/test_registry.ambr +++ /dev/null @@ -1,11 +0,0 @@ -# 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. │ - ╰──────────────────────────────────────────────────────────────────────────────╯ - - ''' -# ---