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

@@ -9,8 +9,10 @@
 * Added convenience function `spcs image-repository url <repo_name>`.
 * 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

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

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")

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

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)

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

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

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"
 )

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"

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"
 )

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."
+    }