From 0d81de051e66e4dc36aa82495e633a36e37bae51 Mon Sep 17 00:00:00 2001 From: David Wang Date: Fri, 9 Feb 2024 05:53:44 -0800 Subject: [PATCH] using show_specific_object in image_repository.manager.get_repository_url (#755) --- .../plugins/spcs/image_repository/manager.py | 33 +------- tests/spcs/test_image_repository.py | 80 ++++++------------- 2 files changed, 28 insertions(+), 85 deletions(-) diff --git a/src/snowflake/cli/plugins/spcs/image_repository/manager.py b/src/snowflake/cli/plugins/spcs/image_repository/manager.py index cfeaddfbbb..5c6e859578 100644 --- a/src/snowflake/cli/plugins/spcs/image_repository/manager.py +++ b/src/snowflake/cli/plugins/spcs/image_repository/manager.py @@ -1,15 +1,12 @@ -from typing import Dict from urllib.parse import urlparse from click import ClickException from snowflake.cli.api.constants import ObjectType from snowflake.cli.api.project.util import ( - escape_like_pattern, is_valid_unquoted_identifier, ) from snowflake.cli.api.sql_execution import SqlExecutionMixin from snowflake.cli.plugins.spcs.common import handle_object_already_exists -from snowflake.connector.cursor import DictCursor from snowflake.connector.errors import ProgrammingError @@ -23,40 +20,16 @@ def get_schema(self): def get_role(self): return self._conn.role - def get_repository_row(self, repo_name: str) -> Dict: + def get_repository_url(self, repo_name: str, with_scheme: bool = True): if not is_valid_unquoted_identifier(repo_name): raise ValueError( f"repo_name '{repo_name}' is not a valid unquoted Snowflake identifier" ) - - repo_name = repo_name.upper() - - # because image repositories only support unquoted identifiers, SHOW LIKE should only return one or zero rows - repository_list_query = ( - f"show image repositories like '{escape_like_pattern(repo_name)}'" - ) - - result_set = self._execute_schema_query( - repository_list_query, cursor_class=DictCursor - ) - results = result_set.fetchall() - - if len(results) == 0: + repo_row = self.show_specific_object("image repositories", repo_name) + if repo_row is None: raise ClickException( f"Image repository '{repo_name}' does not exist in database '{self.get_database()}' and schema '{self.get_schema()}' or not authorized." ) - elif len(results) > 1: - raise ClickException( - f"Found more than one image repository with name matching '{repo_name}'. This is unexpected." - ) - return results[0] - - def get_repository_url(self, repo_name: str, with_scheme: bool = True): - if not is_valid_unquoted_identifier(repo_name): - raise ValueError( - f"repo_name '{repo_name}' is not a valid unquoted Snowflake identifier" - ) - repo_row = self.get_repository_row(repo_name) if with_scheme: return f"https://{repo_row['repository_url']}" else: diff --git a/tests/spcs/test_image_repository.py b/tests/spcs/test_image_repository.py index 32a01ece39..02aed41695 100644 --- a/tests/spcs/test_image_repository.py +++ b/tests/spcs/test_image_repository.py @@ -186,77 +186,47 @@ def test_get_repository_url_cli(mock_url, runner): @mock.patch( - "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._execute_schema_query" -) -def test_get_repository_row(mock_execute, mock_cursor): - mock_execute.return_value = mock_cursor( - rows=MOCK_ROWS_DICT, - columns=MOCK_COLUMNS, - ) - result = ImageRepositoryManager().get_repository_row(MOCK_ROWS_DICT[0]["name"]) - assert result == MOCK_ROWS_DICT[0] - - -@mock.patch( - "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._execute_schema_query" + "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.show_specific_object" ) -@mock.patch( - "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._conn" -) -def test_get_repository_row_no_repo_found(mock_conn, mock_execute, mock_cursor): - mock_execute.return_value = mock_cursor( - rows=[], - columns=MOCK_COLUMNS, - ) - mock_conn.database = "DB" - mock_conn.schema = "SCHEMA" - with pytest.raises(ClickException) as expected_error: - ImageRepositoryManager().get_repository_row("IMAGES") - assert ( - "Image repository 'IMAGES' does not exist in database 'DB' and schema 'SCHEMA' or not authorized." - == expected_error.value.message - ) - - -@mock.patch( - "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._execute_schema_query" -) -def test_get_repository_row_more_than_one_repo(mock_execute, mock_cursor): - mock_execute.return_value = mock_cursor( - rows=MOCK_ROWS_DICT + MOCK_ROWS_DICT, - columns=MOCK_COLUMNS, - ) - with pytest.raises(ClickException) as expected_error: - ImageRepositoryManager().get_repository_row("IMAGES") - assert ( - "Found more than one image repository with name matching 'IMAGES'. This is unexpected." - == expected_error.value.message - ) - - -@mock.patch( - "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.get_repository_row" -) -def test_get_repository_url(mock_get_row, mock_cursor): +def test_get_repository_url(mock_get_row): expected_row = MOCK_ROWS_DICT[0] mock_get_row.return_value = expected_row result = ImageRepositoryManager().get_repository_url(repo_name="IMAGES") - mock_get_row.assert_called_once_with("IMAGES") + mock_get_row.assert_called_once_with("image repositories", "IMAGES") assert isinstance(expected_row, Dict) assert "repository_url" in expected_row assert result == f"https://{expected_row['repository_url']}" @mock.patch( - "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.get_repository_row" + "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.show_specific_object" ) -def test_get_repository_url_no_scheme(mock_get_row, mock_cursor): +def test_get_repository_url_no_scheme(mock_get_row): expected_row = MOCK_ROWS_DICT[0] mock_get_row.return_value = expected_row result = ImageRepositoryManager().get_repository_url( repo_name="IMAGES", with_scheme=False ) - mock_get_row.assert_called_once_with("IMAGES") + mock_get_row.assert_called_once_with("image repositories", "IMAGES") assert isinstance(expected_row, Dict) assert "repository_url" in expected_row assert result == expected_row["repository_url"] + + +@mock.patch( + "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._conn" +) +@mock.patch( + "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.show_specific_object" +) +def test_get_repository_url_no_repo_found(mock_get_row, mock_conn): + mock_get_row.return_value = None + mock_conn.database = "DB" + mock_conn.schema = "SCHEMA" + with pytest.raises(ClickException) as e: + ImageRepositoryManager().get_repository_url(repo_name="IMAGES") + assert ( + e.value.message + == "Image repository 'IMAGES' does not exist in database 'DB' and schema 'SCHEMA' or not authorized." + ) + mock_get_row.assert_called_once_with("image repositories", "IMAGES")