SNOW-1043081: Adding support for qualified names for image repositori…

…es. (#823)

* SNOW-1043081: Adding support for qualified image repository names

* SNOW-1043081: Fixing test imports

* SNOW-1043081: Adding tests for getting image repository url without db or schema

RELEASE-NOTES.md

@@ -4,6 +4,7 @@
 
 ## New additions
 * Added support for fully qualified name (`database.schema.name`) in `name` parameter in streamlit project definition
+* Added support for fully qualified image repository names in `spcs image-repository` commands.
 
 ## Fixes and improvements
 * Adding `--image-name` option for image name argument in `spcs image-repository list-tags` for consistency with other commands.

src/snowflake/cli/api/sql_execution.py

@@ -5,9 +5,8 @@
 from functools import cached_property
 from io import StringIO
 from textwrap import dedent
-from typing import Iterable, Optional
+from typing import Iterable, Optional, Tuple
 
-from click import ClickException
 from snowflake.cli.api.cli_global_context import cli_context
 from snowflake.cli.api.exceptions import (
     DatabaseNotProvidedError,
@@ -19,6 +18,7 @@
     unquote_identifier,
 )
 from snowflake.cli.api.utils.cursor import find_first_row
+from snowflake.cli.api.utils.naming_utils import from_qualified_name
 from snowflake.connector.cursor import DictCursor, SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
 
@@ -82,44 +82,27 @@ def use_role(self, new_role: str):
             if is_different_role:
                 self._execute_query(f"use role {prev_role}")
 
-    def _execute_schema_query(self, query: str, **kwargs):
-        self.check_database_and_schema()
-        return self._execute_query(query, **kwargs)
-
-    def check_database_and_schema(self) -> None:
+    def _execute_schema_query(self, query: str, name: Optional[str] = None, **kwargs):
         """
-        Checks if the connection database and schema are set and that they actually exist in Snowflake.
+        Check that a database and schema are provided before executing the query. Useful for operating on schema level objects.
         """
-        self.check_schema_exists(self._conn.database, self._conn.schema)
+        self.check_database_and_schema_provided(name)
+        return self._execute_query(query, **kwargs)
 
-    def check_database_exists(self, database: str) -> None:
+    def check_database_and_schema_provided(self, name: Optional[str] = None) -> None:
         """
-        Checks that database is provided and that it is a valid database in
-        Snowflake. Note that this could fail for a variety of reasons,
-        including not authorized to use database, database doesn't exist,
-        database is not a valid identifier, and more.
+        Checks if a database and schema are provided, either through the connection context or a qualified name.
         """
+        if name:
+            _, schema, database = from_qualified_name(name)
+        else:
+            schema, database = None, None
+        schema = schema or self._conn.schema
+        database = database or self._conn.database
         if not database:
             raise DatabaseNotProvidedError()
-        try:
-            self._execute_query(f"USE DATABASE {database}")
-        except ProgrammingError as e:
-            raise ClickException(f"Exception occurred: {e}.") from e
-
-    def check_schema_exists(self, database: str, schema: str) -> None:
-        """
-        Checks that schema is provided and that it is a valid schema in Snowflake.
-        Note that this could fail for a variety of reasons,
-        including not authorized to use schema, schema doesn't exist,
-        schema is not a valid identifier, and more.
-        """
-        self.check_database_exists(database)
         if not schema:
             raise SchemaNotProvidedError()
-        try:
-            self._execute_query(f"USE {database}.{schema}")
-        except ProgrammingError as e:
-            raise ClickException(f"Exception occurred: {e}.") from e
 
     def to_fully_qualified_name(
         self, name: str, database: Optional[str] = None, schema: Optional[str] = None
@@ -131,9 +114,7 @@ def to_fully_qualified_name(
 
         if not database:
             if not self._conn.database:
-                raise ClickException(
-                    "Default database not specified in connection details."
-                )
+                raise DatabaseNotProvidedError()
             database = self._conn.database
 
         if len(current_parts) == 2:
@@ -150,29 +131,65 @@ def get_name_from_fully_qualified_name(name):
         Returns name of the object from the fully-qualified name.
         Assumes that [name] is in format [[database.]schema.]name
         """
-        return name.split(".")[-1]
+        return from_qualified_name(name)[0]
+
+    @staticmethod
+    def _qualified_name_to_in_clause(name: str) -> Tuple[str, Optional[str]]:
+        unqualified_name, schema, database = from_qualified_name(name)
+        if database:
+            in_clause = f"in schema {database}.{schema}"
+        elif schema:
+            in_clause = f"in schema {schema}"
+        else:
+            in_clause = None
+        return unqualified_name, in_clause
+
+    class InClauseWithQualifiedNameError(ValueError):
+        def __init__(self):
+            super().__init__("non-empty 'in_clause' passed with qualified 'name'")
 
     def show_specific_object(
         self,
         object_type_plural: str,
-        unqualified_name: str,
+        name: str,
         name_col: str = "name",
         in_clause: str = "",
         check_schema: bool = False,
     ) -> Optional[dict]:
         """
         Executes a "show <objects> like" query for a particular entity with a
-        given (unqualified) name. This command is useful when the corresponding
+        given (optionally qualified) name. This command is useful when the corresponding
         "describe <object>" query does not provide the information you seek.
+
+        Note that this command is analogous to describe and should only return a single row.
+        If the target object type is a schema level object, then check_schema should be set to True
+        so that the function will verify that a database and schema are provided, either through
+        the connection or a qualified name, before executing the query.
         """
-        if check_schema:
-            self.check_database_and_schema()
+
+        unqualified_name, name_in_clause = self._qualified_name_to_in_clause(name)
+        if in_clause and name_in_clause:
+            raise self.InClauseWithQualifiedNameError()
+        elif name_in_clause:
+            in_clause = name_in_clause
         show_obj_query = f"show {object_type_plural} like {identifier_to_show_like_pattern(unqualified_name)} {in_clause}".strip()
-        show_obj_cursor = self._execute_query(  # type: ignore
-            show_obj_query, cursor_class=DictCursor
-        )
+
+        if check_schema:
+            show_obj_cursor = self._execute_schema_query(  # type: ignore
+                show_obj_query, name=name, cursor_class=DictCursor
+            )
+        else:
+            show_obj_cursor = self._execute_query(  # type: ignore
+                show_obj_query, cursor_class=DictCursor
+            )
+
         if show_obj_cursor.rowcount is None:
             raise SnowflakeSQLExecutionError(show_obj_query)
+        elif show_obj_cursor.rowcount > 1:
+            raise ProgrammingError(
+                f"Received multiple rows from result of SQL statement: {show_obj_query}. Usage of 'show_specific_object' may not be properly scoped."
+            )
+
         show_obj_row = find_first_row(
             show_obj_cursor,
             lambda row: row[name_col] == unquote_identifier(unqualified_name),

src/snowflake/cli/api/utils/naming_utils.py

@@ -0,0 +1,27 @@
+import re
+from typing import Optional, Tuple
+
+from snowflake.cli.api.project.util import (
+    VALID_IDENTIFIER_REGEX,
+)
+
+
+def from_qualified_name(name: str) -> Tuple[str, Optional[str], Optional[str]]:
+    """
+    Takes in an object name in the form [[database.]schema.]name. Returns a tuple (name, [schema], [database])
+    """
+    # TODO: Use regex to match object name to a valid identifier or valid identifier (args). Second case is for sprocs and UDFs
+    qualifier_pattern = rf"(?:(?P<first_qualifier>{VALID_IDENTIFIER_REGEX})\.)?(?:(?P<second_qualifier>{VALID_IDENTIFIER_REGEX})\.)?(?P<name>.*)"
+    result = re.fullmatch(qualifier_pattern, name)
+
+    if result is None:
+        raise ValueError(f"'{name}' is not a valid qualified name")
+
+    unqualified_name = result.group("name")
+    if result.group("second_qualifier") is not None:
+        database = result.group("first_qualifier")
+        schema = result.group("second_qualifier")
+    else:
+        database = None
+        schema = result.group("first_qualifier")
+    return unqualified_name, schema, database

src/snowflake/cli/plugins/object/common.py

@@ -4,8 +4,7 @@
 from click import ClickException
 from snowflake.cli.api.commands.flags import OverrideableOption
 from snowflake.cli.api.project.util import (
-    QUOTED_IDENTIFIER_REGEX,
-    UNQUOTED_IDENTIFIER_REGEX,
+    VALID_IDENTIFIER_REGEX,
     is_valid_identifier,
     to_string_literal,
 )
@@ -34,11 +33,9 @@ def __init__(self):
 def _parse_tag(tag: str) -> Tag:
     import re
 
-    identifier_pattern = re.compile(
-        f"(?P<tag_name>{UNQUOTED_IDENTIFIER_REGEX}|{QUOTED_IDENTIFIER_REGEX})"
-    )
-    value_pattern = re.compile(f"(?P<tag_value>.+)")
-    result = re.fullmatch(f"{identifier_pattern.pattern}={value_pattern.pattern}", tag)
+    identifier_pattern = rf"(?P<tag_name>{VALID_IDENTIFIER_REGEX})"
+    value_pattern = r"(?P<tag_value>.+)"
+    result = re.fullmatch(rf"{identifier_pattern}={value_pattern}", tag)
     if result is not None:
         try:
             return Tag(result.group("tag_name"), result.group("tag_value"))

src/snowflake/cli/plugins/spcs/image_repository/commands.py

@@ -22,7 +22,7 @@
 
 
 def _repo_name_callback(name: str):
-    if not is_valid_object_name(name, max_depth=0, allow_quoted=True):
+    if not is_valid_object_name(name, max_depth=2, allow_quoted=False):
         raise ClickException(
             f"'{name}' is not a valid image repository name. Note that image repository names must be unquoted identifiers. The same constraint also applies to database and schema names where you create an image repository."
         )

src/snowflake/cli/plugins/spcs/image_repository/manager.py

@@ -1,10 +1,6 @@
 from urllib.parse import urlparse
 
-from click import ClickException
 from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.api.project.util import (
-    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.errors import ProgrammingError
@@ -21,17 +17,13 @@ def get_role(self):
         return self._conn.role
 
     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"
-            )
-        # we explicitly do not allow this function to be used without connection database and schema set
+
         repo_row = self.show_specific_object(
             "image repositories", repo_name, check_schema=True
         )
         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."
+            raise ProgrammingError(
+                f"Image repository '{self.to_fully_qualified_name(repo_name)}' does not exist or not authorized."
             )
         if with_scheme:
             return f"https://{repo_row['repository_url']}"
@@ -53,6 +45,8 @@ def get_repository_api_url(self, repo_url):
 
     def create(self, name: str):
         try:
-            return self._execute_schema_query(f"create image repository {name}")
+            return self._execute_schema_query(
+                f"create image repository {name}", name=name
+            )
         except ProgrammingError as e:
             handle_object_already_exists(e, ObjectType.IMAGE_REPOSITORY, name)

tests/api/utils/test_naming_utils.py

@@ -0,0 +1,15 @@
+import pytest
+from snowflake.cli.api.utils.naming_utils import from_qualified_name
+
+
+@pytest.mark.parametrize(
+    "qualified_name, expected",
+    [
+        ("func(number, number)", ("func(number, number)", None, None)),
+        ("name", ("name", None, None)),
+        ("schema.name", ("name", "schema", None)),
+        ("db.schema.name", ("name", "schema", "db")),
+    ],
+)
+def test_from_fully_qualified_name(qualified_name, expected):
+    assert from_qualified_name(qualified_name) == expected

tests/spcs/test_image_repository.py

@@ -4,14 +4,14 @@
 from unittest.mock import Mock
 
 import pytest
-from click import ClickException
 from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.exceptions import (
     DatabaseNotProvidedError,
     SchemaNotProvidedError,
 )
 from snowflake.cli.plugins.spcs.image_repository.manager import ImageRepositoryManager
 from snowflake.connector.cursor import SnowflakeCursor
+from snowflake.connector.errors import ProgrammingError
 
 from tests.spcs.test_common import SPCS_OBJECT_EXISTS_ERROR
 
@@ -55,7 +55,7 @@ def test_create(
     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)
+    mock_execute.assert_called_once_with(expected_query, name=repo_name)
     assert result == cursor
 
 
@@ -191,13 +191,10 @@ def test_get_repository_url_cli(mock_url, runner):
     assert result.output.strip() == repo_url
 
 
-@mock.patch(
-    "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.check_database_and_schema"
-)
 @mock.patch(
     "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.show_specific_object"
 )
-def test_get_repository_url(mock_get_row, mock_check_database_and_schema):
+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")
@@ -209,13 +206,10 @@ def test_get_repository_url(mock_get_row, mock_check_database_and_schema):
     assert result == f"https://{expected_row['repository_url']}"
 
 
-@mock.patch(
-    "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.check_database_and_schema"
-)
 @mock.patch(
     "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.show_specific_object"
 )
-def test_get_repository_url_no_scheme(mock_get_row, mock_check_database_and_schema):
+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(
@@ -229,26 +223,21 @@ def test_get_repository_url_no_scheme(mock_get_row, mock_check_database_and_sche
     assert result == expected_row["repository_url"]
 
 
-@mock.patch(
-    "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager.check_database_and_schema"
-)
 @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_check_database_and_schema
-):
+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:
+    with pytest.raises(ProgrammingError) 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."
+        e.value.msg
+        == "Image repository 'DB.SCHEMA.IMAGES' does not exist or not authorized."
     )
     mock_get_row.assert_called_once_with(
         "image repositories", "IMAGES", check_schema=True
@@ -258,17 +247,17 @@ def test_get_repository_url_no_repo_found(
 @mock.patch(
     "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._conn"
 )
-def test_get_repository_url_no_database(mock_conn):
+def test_get_repository_url_no_database_provided(mock_conn):
     mock_conn.database = None
     with pytest.raises(DatabaseNotProvidedError):
-        ImageRepositoryManager().get_repository_url("test_repo")
+        ImageRepositoryManager().get_repository_url("IMAGES")
 
 
 @mock.patch(
     "snowflake.cli.plugins.spcs.image_repository.manager.ImageRepositoryManager._conn"
 )
-@mock.patch("snowflake.cli.api.sql_execution.SqlExecutionMixin.check_database_exists")
-def test_get_repository_url_no_schema(mock_check_database_exists, mock_conn):
+def test_get_repository_url_no_schema_provided(mock_conn):
+    mock_conn.database = "DB"
     mock_conn.schema = None
     with pytest.raises(SchemaNotProvidedError):
-        ImageRepositoryManager().get_repository_url("test_repo")
+        ImageRepositoryManager().get_repository_url("IMAGES")