From aaa5889a9d530cd50217c9f330d78669e62e94cb Mon Sep 17 00:00:00 2001 From: David Wang Date: Wed, 7 Feb 2024 09:40:43 -0800 Subject: [PATCH] SNOW-1011761: Adding `image-repository create` command (#746) * SNOW-1011761: Added create image repository command. * SNOW-1011761: Adding more informative errors for when compute pool or image repo already exists and you try to create one * SNOW-1011761: Update release notes * SNOW-1011761: Bug fix to remove IF NOT EXISTS from 'spcs service create' * SNOW-1011761: Typo fix * fixup * SNOW-1011761: Factoring our common error handling for ObjectAlreadyExistsError between services, compute pools, and image repositories * fixup * SNOW-1011761: Adding unit test for case where image repository exists * SNOW-1011761: Adding unit tests for logic when creating an image repository, service, or compute pool that already exists * linting --- RELEASE-NOTES.md | 2 + src/snowflake/cli/plugins/spcs/common.py | 16 ++++++ .../cli/plugins/spcs/compute_pool/manager.py | 13 ++++- .../plugins/spcs/image_repository/commands.py | 19 ++++++- .../plugins/spcs/image_repository/manager.py | 9 +++ .../cli/plugins/spcs/services/manager.py | 14 ++++- tests/spcs/test_common.py | 24 ++++++++ tests/spcs/test_compute_pool.py | 24 ++++++++ tests/spcs/test_image_repository.py | 55 ++++++++++++++++++- tests/spcs/test_services.py | 39 +++++++++++-- .../spcs/test_image_repository.py | 14 ++++- 11 files changed, 215 insertions(+), 14 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 6b163dbfe5..13afe4be51 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -9,8 +9,10 @@ * Added convenience function `spcs image-repository url `. * Added `suspend` and `resume` commands for `spcs compute-pool`. +* Added `create` command to `spcs image-repository`. ## Fixes and improvements * Restricted permissions of automatically created files +* Fixed bug where `spcs service create` would not throw error if service with specified name already exists. # v2.0.0 diff --git a/src/snowflake/cli/plugins/spcs/common.py b/src/snowflake/cli/plugins/spcs/common.py index e4fdfc5e74..1bd635be02 100644 --- a/src/snowflake/cli/plugins/spcs/common.py +++ b/src/snowflake/cli/plugins/spcs/common.py @@ -4,6 +4,10 @@ from typing import TextIO from click import ClickException +from snowflake.cli.api.constants import ObjectType +from snowflake.cli.api.exceptions import ObjectAlreadyExistsError +from snowflake.cli.api.project.util import unquote_identifier +from snowflake.connector.errors import ProgrammingError if not sys.stdout.closed and sys.stdout.isatty(): GREEN = "\033[32m" @@ -59,3 +63,15 @@ def validate_and_set_instances(min_instances, max_instances, instance_name): f"max_{instance_name} must be greater or equal to min_{instance_name}" ) return max_instances + + +def handle_object_already_exists( + error: ProgrammingError, object_type: ObjectType, object_name: str +): + if error.errno == 2002: + raise ObjectAlreadyExistsError( + object_type=object_type, + name=unquote_identifier(object_name), + ) + else: + raise error diff --git a/src/snowflake/cli/plugins/spcs/compute_pool/manager.py b/src/snowflake/cli/plugins/spcs/compute_pool/manager.py index 897cd0ba47..d112927b83 100644 --- a/src/snowflake/cli/plugins/spcs/compute_pool/manager.py +++ b/src/snowflake/cli/plugins/spcs/compute_pool/manager.py @@ -1,8 +1,13 @@ from typing import Optional +from snowflake.cli.api.constants import ObjectType from snowflake.cli.api.sql_execution import SqlExecutionMixin -from snowflake.cli.plugins.spcs.common import strip_empty_lines +from snowflake.cli.plugins.spcs.common import ( + handle_object_already_exists, + strip_empty_lines, +) from snowflake.connector.cursor import SnowflakeCursor +from snowflake.connector.errors import ProgrammingError class ComputePoolManager(SqlExecutionMixin): @@ -28,7 +33,11 @@ def create( """.splitlines() if comment: query.append(f"COMMENT = {comment}") - return self._execute_query(strip_empty_lines(query)) + + try: + return self._execute_query(strip_empty_lines(query)) + except ProgrammingError as e: + handle_object_already_exists(e, ObjectType.COMPUTE_POOL, pool_name) def stop(self, pool_name: str) -> SnowflakeCursor: return self._execute_query(f"alter compute pool {pool_name} stop all") diff --git a/src/snowflake/cli/plugins/spcs/image_repository/commands.py b/src/snowflake/cli/plugins/spcs/image_repository/commands.py index 45384969c4..ccbaec93af 100644 --- a/src/snowflake/cli/plugins/spcs/image_repository/commands.py +++ b/src/snowflake/cli/plugins/spcs/image_repository/commands.py @@ -9,7 +9,11 @@ with_output, ) from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS -from snowflake.cli.api.output.types import CollectionResult, MessageResult +from snowflake.cli.api.output.types import ( + CollectionResult, + MessageResult, + SingleQueryResult, +) from snowflake.cli.api.project.util import is_valid_unquoted_identifier from snowflake.cli.plugins.spcs.image_registry.manager import RegistryManager from snowflake.cli.plugins.spcs.image_repository.manager import ImageRepositoryManager @@ -36,6 +40,19 @@ def _repo_name_callback(name: str): ) +@app.command() +@with_output +@global_options_with_connection +def create( + name: str = REPO_NAME_ARGUMENT, + **options, +): + """ + Creates a new image repository in the current schema. + """ + return SingleQueryResult(ImageRepositoryManager().create(name=name)) + + @app.command("list-images") @with_output @global_options_with_connection diff --git a/src/snowflake/cli/plugins/spcs/image_repository/manager.py b/src/snowflake/cli/plugins/spcs/image_repository/manager.py index c226e2618c..cfeaddfbbb 100644 --- a/src/snowflake/cli/plugins/spcs/image_repository/manager.py +++ b/src/snowflake/cli/plugins/spcs/image_repository/manager.py @@ -2,12 +2,15 @@ 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 class ImageRepositoryManager(SqlExecutionMixin): @@ -71,3 +74,9 @@ def get_repository_api_url(self, repo_url): path = parsed_url.path return f"{scheme}://{host}/v2{path}" + + def create(self, name: str): + try: + return self._execute_schema_query(f"create image repository {name}") + except ProgrammingError as e: + handle_object_already_exists(e, ObjectType.IMAGE_REPOSITORY, name) diff --git a/src/snowflake/cli/plugins/spcs/services/manager.py b/src/snowflake/cli/plugins/spcs/services/manager.py index fc24fb26b5..21b6947d79 100644 --- a/src/snowflake/cli/plugins/spcs/services/manager.py +++ b/src/snowflake/cli/plugins/spcs/services/manager.py @@ -1,10 +1,15 @@ from pathlib import Path from typing import List, Optional +from snowflake.cli.api.constants import ObjectType from snowflake.cli.api.sql_execution import SqlExecutionMixin from snowflake.cli.plugins.object.common import Tag -from snowflake.cli.plugins.spcs.common import strip_empty_lines +from snowflake.cli.plugins.spcs.common import ( + handle_object_already_exists, + strip_empty_lines, +) from snowflake.connector.cursor import SnowflakeCursor +from snowflake.connector.errors import ProgrammingError class ServiceManager(SqlExecutionMixin): @@ -24,7 +29,7 @@ def create( spec = self._read_yaml(spec_path) query = f"""\ - CREATE SERVICE IF NOT EXISTS {service_name} + CREATE SERVICE {service_name} IN COMPUTE POOL {compute_pool} FROM SPECIFICATION $$ {spec} @@ -53,7 +58,10 @@ def create( if comment: query.append(f"COMMENT = {comment}") - return self._execute_schema_query(strip_empty_lines(query)) + try: + return self._execute_schema_query(strip_empty_lines(query)) + except ProgrammingError as e: + handle_object_already_exists(e, ObjectType.SERVICE, service_name) def _read_yaml(self, path: Path) -> str: # TODO(aivanou): Add validation towards schema diff --git a/tests/spcs/test_common.py b/tests/spcs/test_common.py index a6702edadd..29dd04a258 100644 --- a/tests/spcs/test_common.py +++ b/tests/spcs/test_common.py @@ -1,6 +1,10 @@ from snowflake.cli.plugins.spcs.common import validate_and_set_instances from tests.testing_utils.fixtures import * from click import ClickException +from snowflake.connector.errors import ProgrammingError +from snowflake.cli.plugins.spcs.common import handle_object_already_exists +from snowflake.cli.api.exceptions import ObjectAlreadyExistsError, ObjectType +from unittest.mock import Mock @pytest.mark.parametrize( @@ -36,3 +40,23 @@ def test_validate_and_set_instances_invalid(min_instances, max_instances, expect with pytest.raises(ClickException) as exc: validate_and_set_instances(min_instances, max_instances, "name") assert expected_msg in exc.value.message + + +SPCS_OBJECT_EXISTS_ERROR = ProgrammingError( + msg="Object 'TEST_OBJECT' already exists.", errno=2002 +) + + +def test_handle_object_exists_error(): + mock_type = Mock(spec=ObjectType) + test_name = "TEST_OBJECT" + with pytest.raises(ObjectAlreadyExistsError): + handle_object_already_exists(SPCS_OBJECT_EXISTS_ERROR, mock_type, test_name) + + +def test_handle_object_exists_error_other_error(): + # For any errors other than 'Object 'XYZ' already exists.', simply pass the error through + other_error = ProgrammingError(msg="Object does not already exist.", errno=0) + with pytest.raises(ProgrammingError) as e: + handle_object_already_exists(other_error, Mock(spec=ObjectType), "TEST_OBJECT") + assert other_error == e.value diff --git a/tests/spcs/test_compute_pool.py b/tests/spcs/test_compute_pool.py index 80c06aba97..9eaa0c1e6c 100644 --- a/tests/spcs/test_compute_pool.py +++ b/tests/spcs/test_compute_pool.py @@ -8,6 +8,8 @@ import pytest from click import ClickException +from tests.spcs.test_common import SPCS_OBJECT_EXISTS_ERROR +from snowflake.cli.api.constants import ObjectType @patch( @@ -111,6 +113,28 @@ def test_create_pool_cli(mock_create, runner): ) +@patch( + "snowflake.cli.plugins.spcs.compute_pool.manager.ComputePoolManager._execute_query" +) +@patch("snowflake.cli.plugins.spcs.compute_pool.manager.handle_object_already_exists") +def test_create_repository_already_exists(mock_handle, mock_execute): + pool_name = "test_object" + mock_execute.side_effect = SPCS_OBJECT_EXISTS_ERROR + ComputePoolManager().create( + pool_name=pool_name, + min_nodes=1, + max_nodes=1, + instance_family="test_family", + auto_resume=False, + initially_suspended=True, + auto_suspend_secs=7200, + comment=to_string_literal("this is a test"), + ) + mock_handle.assert_called_once_with( + SPCS_OBJECT_EXISTS_ERROR, ObjectType.COMPUTE_POOL, pool_name + ) + + @patch( "snowflake.cli.plugins.spcs.compute_pool.manager.ComputePoolManager._execute_query" ) diff --git a/tests/spcs/test_image_repository.py b/tests/spcs/test_image_repository.py index 335a48ae39..32a01ece39 100644 --- a/tests/spcs/test_image_repository.py +++ b/tests/spcs/test_image_repository.py @@ -1,10 +1,13 @@ -import snowflake.cli.plugins.spcs.image_repository.manager from tests.testing_utils.fixtures import * import json from snowflake.cli.plugins.spcs.image_repository.manager import ImageRepositoryManager from typing import Dict from click import ClickException - +from unittest.mock import Mock +from snowflake.connector.cursor import SnowflakeCursor +from snowflake.connector.errors import ProgrammingError +from snowflake.cli.api.constants import ObjectType +from tests.spcs.test_common import SPCS_OBJECT_EXISTS_ERROR MOCK_ROWS = [ [ @@ -35,6 +38,54 @@ ] +@mock.patch( + "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._execute_schema_query" +) +def test_create( + mock_execute, +): + repo_name = "test_repo" + cursor = Mock(spec=SnowflakeCursor) + mock_execute.return_value = cursor + result = ImageRepositoryManager().create(name=repo_name) + expected_query = "create image repository test_repo" + mock_execute.assert_called_once_with(expected_query) + assert result == cursor + + +@mock.patch( + "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.create" +) +def test_create_cli(mock_create, mock_cursor, runner): + repo_name = "test_repo" + cursor = mock_cursor( + rows=[[f"Image Repository {repo_name.upper()} successfully created."]], + columns=["status"], + ) + mock_create.return_value = cursor + result = runner.invoke(["spcs", "image-repository", "create", repo_name]) + mock_create.assert_called_once_with(name=repo_name) + assert result.exit_code == 0, result.output + assert ( + f"Image Repository {repo_name.upper()} successfully created." in result.output + ) + + +@mock.patch( + "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._execute_schema_query" +) +@mock.patch( + "snowflake.cli.plugins.spcs.image_repository.manager.handle_object_already_exists" +) +def test_create_repository_already_exists(mock_handle, mock_execute): + repo_name = "test_object" + mock_execute.side_effect = SPCS_OBJECT_EXISTS_ERROR + ImageRepositoryManager().create(repo_name) + mock_handle.assert_called_once_with( + SPCS_OBJECT_EXISTS_ERROR, ObjectType.IMAGE_REPOSITORY, repo_name + ) + + @mock.patch("snowflake.cli.plugins.spcs.image_repository.commands.requests.get") @mock.patch( "snowflake.cli.plugins.spcs.image_repository.commands.ImageRepositoryManager._execute_query" diff --git a/tests/spcs/test_services.py b/tests/spcs/test_services.py index bb4f989873..dbe248ba78 100644 --- a/tests/spcs/test_services.py +++ b/tests/spcs/test_services.py @@ -1,10 +1,8 @@ -from pathlib import Path from textwrap import dedent from unittest.mock import Mock, patch -from click import ClickException -import pytest -import strictyaml +from tests.spcs.test_common import SPCS_OBJECT_EXISTS_ERROR +from snowflake.cli.api.constants import ObjectType from snowflake.cli.plugins.spcs.services.manager import ServiceManager from tests.testing_utils.fixtures import * from snowflake.cli.api.project.util import to_string_literal @@ -62,7 +60,7 @@ def test_create_service(mock_execute_schema_query, other_directory): ) expected_query = " ".join( [ - "CREATE SERVICE IF NOT EXISTS test_service", + "CREATE SERVICE test_service", "IN COMPUTE POOL test_pool", 'FROM SPECIFICATION $$ {"spec": {"containers": [{"name": "cloudbeaver", "image":', '"/spcs_demos_db/cloudbeaver:23.2.1"}], "endpoints": [{"name": "cloudbeaver",', @@ -216,6 +214,37 @@ def test_create_service_with_invalid_spec(mock_read_yaml): ) +@patch("snowflake.cli.plugins.spcs.services.manager.ServiceManager._read_yaml") +@patch( + "snowflake.cli.plugins.spcs.services.manager.ServiceManager._execute_schema_query" +) +@patch("snowflake.cli.plugins.spcs.services.manager.handle_object_already_exists") +def test_create_repository_already_exists(mock_handle, mock_execute, mock_read_yaml): + service_name = "test_object" + compute_pool = "test_pool" + spec_path = "/path/to/spec.yaml" + min_instances = 42 + max_instances = 42 + external_access_integrations = query_warehouse = tags = comment = None + auto_resume = False + mock_execute.side_effect = SPCS_OBJECT_EXISTS_ERROR + ServiceManager().create( + service_name=service_name, + compute_pool=compute_pool, + spec_path=Path(spec_path), + min_instances=min_instances, + max_instances=max_instances, + auto_resume=auto_resume, + external_access_integrations=external_access_integrations, + query_warehouse=query_warehouse, + tags=tags, + comment=comment, + ) + mock_handle.assert_called_once_with( + SPCS_OBJECT_EXISTS_ERROR, ObjectType.SERVICE, service_name + ) + + @patch( "snowflake.cli.plugins.spcs.services.manager.ServiceManager._execute_schema_query" ) diff --git a/tests_integration/spcs/test_image_repository.py b/tests_integration/spcs/test_image_repository.py index 580f8c8b36..3d82385d07 100644 --- a/tests_integration/spcs/test_image_repository.py +++ b/tests_integration/spcs/test_image_repository.py @@ -2,7 +2,7 @@ from snowflake.cli.api.project.util import escape_like_pattern from tests_integration.test_utils import contains_row_with, row_from_snowflake_session -from tests_integration.testing_utils.naming_utils import ObjectNameProvider +from tests_integration.testing_utils import ObjectNameProvider INTEGRATION_DATABASE = "SNOWCLI_DB" INTEGRATION_SCHEMA = "PUBLIC" @@ -80,3 +80,15 @@ def test_get_repo_url(runner, snowflake_session, test_database): ) assert isinstance(result.output, str), result.output assert result.output.strip() == expect_url + + +@pytest.mark.integration +def test_create_image_repository(runner, test_database): + repo_name = ObjectNameProvider("test_repo").create_and_get_next_object_name() + result = runner.invoke_with_connection_json( + ["spcs", "image-repository", "create", repo_name] + ) + assert isinstance(result.json, dict), result.output + assert result.json == { + "status": f"Image Repository {repo_name.upper()} successfully created." + }