Skip to content

Commit

Permalink
using show_specific_object in image_repository.manager.get_repository…
Browse files Browse the repository at this point in the history 
…_url (#755)
  • Loading branch information
@sfc-gh-davwang @sfc-gh-cgorrie
sfc-gh-davwang authored and sfc-gh-cgorrie committed Feb 9, 2024
1 parent 6f16c7b commit 0d81de0
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 85 deletions.
33 changes: 3 additions & 30 deletions src/snowflake/cli/plugins/spcs/image_repository/manager.py
View file
Original file line number Diff line number Diff line change
@@ -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


Expand All @@ -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:
Expand Down
80 changes: 25 additions & 55 deletions tests/spcs/test_image_repository.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")

0 comments on commit 0d81de0

Please sign in to comment.