snowflake-cli-labs 3.0.0rc1__py3-none-any.whl → 3.0.0rc3__py3-none-any.whl

snowflake/cli/_plugins/nativeapp/version/version_processor.py

@@ -16,17 +16,17 @@ from __future__ import annotations
 
 from pathlib import Path
 from textwrap import dedent
-from typing import Dict, List, Optional
+from typing import Dict, Optional
 
 import typer
-from click import BadOptionUsage, ClickException
+from click import ClickException
+from snowflake.cli._plugins.nativeapp.application_package_entity import (
+    ApplicationPackageEntity,
+)
 from snowflake.cli._plugins.nativeapp.artifacts import (
-    BundleMap,
     find_version_info_in_manifest_file,
 )
-from snowflake.cli._plugins.nativeapp.constants import VERSION_COL
 from snowflake.cli._plugins.nativeapp.exceptions import (
-    ApplicationPackageAlreadyExistsError,
     ApplicationPackageDoesNotExistError,
 )
 from snowflake.cli._plugins.nativeapp.manager import (
@@ -36,245 +36,50 @@ from snowflake.cli._plugins.nativeapp.manager import (
 from snowflake.cli._plugins.nativeapp.policy import PolicyBase
 from snowflake.cli._plugins.nativeapp.run_processor import NativeAppRunProcessor
 from snowflake.cli.api.console import cli_console as cc
-from snowflake.cli.api.entities.utils import ensure_correct_owner
-from snowflake.cli.api.exceptions import SnowflakeSQLExecutionError
-from snowflake.cli.api.project.schemas.native_app.native_app import NativeApp
-from snowflake.cli.api.project.util import to_identifier, unquote_identifier
-from snowflake.cli.api.utils.cursor import (
-    find_all_rows,
-)
+from snowflake.cli.api.project.schemas.v1.native_app.native_app import NativeApp
+from snowflake.cli.api.project.util import to_identifier
 from snowflake.connector import ProgrammingError
-from snowflake.connector.cursor import DictCursor
-
-
-def check_index_changes_in_git_repo(
-    project_root: Path, policy: PolicyBase, is_interactive: bool
-) -> None:
-    """
-    Checks if the project root, i.e. the native apps project is a git repository. If it is a git repository,
-    it also checks if there any local changes to the directory that may not be on the application package stage.
-    """
-    from git import Repo
-    from git.exc import InvalidGitRepositoryError
-
-    try:
-        repo = Repo(project_root, search_parent_directories=True)
-        assert repo.git_dir is not None
-
-        # Check if the repo has any changes, including untracked files
-        if repo.is_dirty(untracked_files=True):
-            cc.warning(
-                "Changes detected in the git repository. "
-                "(Rerun your command with --skip-git-check flag to ignore this check)"
-            )
-            repo.git.execute(["git", "status"])
-
-            user_prompt = (
-                "You have local changes in this repository that are not part of a previous commit. "
-                "Do you still want to continue?"
-            )
-            if not policy.should_proceed(user_prompt):
-                if is_interactive:
-                    cc.message("Not creating a new version.")
-                    raise typer.Exit(0)
-                else:
-                    cc.message(
-                        "Cannot create a new version non-interactively without --force."
-                    )
-                    raise typer.Exit(1)
-
-    except InvalidGitRepositoryError:
-        pass  # not a git repository, which is acceptable
 
 
 class NativeAppVersionCreateProcessor(NativeAppRunProcessor):
     def __init__(self, project_definition: Dict, project_root: Path):
         super().__init__(project_definition, project_root)
 
-    def get_existing_release_directive_info_for_version(
-        self, version: str
-    ) -> List[dict]:
-        """
-        Get all existing release directives, if present, set on the version defined in an application package.
-        It executes a 'show release directives in application package' query and returns the filtered results, if they exist.
-        """
-        with self.use_role(self.package_role):
-            show_obj_query = (
-                f"show release directives in application package {self.package_name}"
-            )
-            show_obj_cursor = self._execute_query(
-                show_obj_query, cursor_class=DictCursor
-            )
-
-            if show_obj_cursor.rowcount is None:
-                raise SnowflakeSQLExecutionError(show_obj_query)
-
-            show_obj_rows = find_all_rows(
-                show_obj_cursor,
-                lambda row: row[VERSION_COL] == unquote_identifier(version),
-            )
-
-            return show_obj_rows
-
-    def add_new_version(self, version: str) -> None:
-        """
-        Defines a new version in an existing application package.
-        """
-        # Make the version a valid identifier, adding quotes if necessary
-        version = to_identifier(version)
-        with self.use_role(self.package_role):
-            cc.step(
-                f"Defining a new version {version} in application package {self.package_name}"
-            )
-            add_version_query = dedent(
-                f"""\
-                    alter application package {self.package_name}
-                        add version {version}
-                        using @{self.stage_fqn}
-                """
-            )
-            self._execute_query(add_version_query, cursor_class=DictCursor)
-            cc.message(
-                f"Version {version} created for application package {self.package_name}."
-            )
-
-    def add_new_patch_to_version(self, version: str, patch: Optional[int] = None):
-        """
-        Add a new patch, optionally a custom one, to an existing version in an application package.
-        """
-        # Make the version a valid identifier, adding quotes if necessary
-        version = to_identifier(version)
-        with self.use_role(self.package_role):
-            cc.step(
-                f"Adding new patch to version {version} defined in application package {self.package_name}"
-            )
-            add_version_query = dedent(
-                f"""\
-                    alter application package {self.package_name}
-                        add patch {patch if patch else ""} for version {version}
-                        using @{self.stage_fqn}
-                """
-            )
-            result_cursor = self._execute_query(
-                add_version_query, cursor_class=DictCursor
-            )
-
-            show_row = result_cursor.fetchall()[0]
-            new_patch = show_row["patch"]
-            cc.message(
-                f"Patch {new_patch} created for version {version} defined in application package {self.package_name}."
-            )
-
     def process(
         self,
-        bundle_map: BundleMap,
         version: Optional[str],
         patch: Optional[int],
-        policy: PolicyBase,
-        git_policy: PolicyBase,
-        is_interactive: bool,
+        force: bool,
+        interactive: bool,
+        skip_git_check: bool,
         *args,
         **kwargs,
     ):
-        """
-        Perform bundle, application package creation, stage upload, version and/or patch to an application package.
-        """
-
-        # Make sure version is not None before proceeding any further.
-        # This will raise an exception if version information is not found. Patch can be None.
-        if not version:
-            cc.message(
-                "Version was not provided through the Snowflake CLI. Checking version in the manifest.yml instead."
-            )
-
-            version, patch = find_version_info_in_manifest_file(self.deploy_root)
-            if not version:
-                raise ClickException(
-                    "Manifest.yml file does not contain a value for the version field."
-                )
-
-        # Check if --patch needs to throw a bad option error, either if application package does not exist or if version does not exist
-        if patch:
-            try:
-                if not self.get_existing_version_info(version):
-                    raise BadOptionUsage(
-                        option_name="patch",
-                        message=f"Cannot create a custom patch when version {version} is not defined in the application package {self.package_name}. Try again without using --patch.",
-                    )
-            except ApplicationPackageDoesNotExistError as app_err:
-                raise BadOptionUsage(
-                    option_name="patch",
-                    message=f"Cannot create a custom patch when application package {self.package_name} does not exist. Try again without using --patch.",
-                )
-
-        if git_policy.should_proceed():
-            check_index_changes_in_git_repo(
-                project_root=self.project_root,
-                policy=policy,
-                is_interactive=is_interactive,
-            )
-
-        # TODO: consider using self.deploy() instead
-
-        try:
-            self.create_app_package()
-        except ApplicationPackageAlreadyExistsError as e:
-            cc.warning(e.message)
-            if not policy.should_proceed("Proceed with using this package?"):
-                raise typer.Abort() from e
-
-        with self.use_role(self.package_role):
-            # Now that the application package exists, create shared data
-            self._apply_package_scripts()
-
-            # Upload files from deploy root local folder to the above stage
-            self.sync_deploy_root_with_stage(
-                bundle_map=bundle_map,
-                role=self.package_role,
-                prune=True,
-                recursive=True,
-                stage_fqn=self.stage_fqn,
-            )
-            with self.use_package_warehouse():
-                self.execute_package_post_deploy_hooks()
-
-        # Warn if the version exists in a release directive(s)
-        existing_release_directives = (
-            self.get_existing_release_directive_info_for_version(version)
+        return ApplicationPackageEntity.version_create(
+            console=cc,
+            project_root=self.project_root,
+            deploy_root=self.deploy_root,
+            bundle_root=self.bundle_root,
+            generated_root=self.generated_root,
+            artifacts=self.artifacts,
+            package_name=self.package_name,
+            package_role=self.package_role,
+            package_distribution=self.package_distribution,
+            prune=True,
+            recursive=True,
+            paths=None,
+            print_diff=True,
+            validate=True,
+            stage_fqn=self.stage_fqn,
+            package_warehouse=self.package_warehouse,
+            post_deploy_hooks=self.package_post_deploy_hooks,
+            package_scripts=self.package_scripts,
+            version=version,
+            patch=patch,
+            force=force,
+            interactive=interactive,
+            skip_git_check=skip_git_check,
         )
-        if existing_release_directives:
-            release_directive_names = ", ".join(
-                row["name"] for row in existing_release_directives
-            )
-            cc.warning(
-                dedent(
-                    f"""\
-                    Version {version} already defined in application package {self.package_name} and in release directive(s): {release_directive_names}.
-                    """
-                )
-            )
-
-            user_prompt = (
-                f"Are you sure you want to create a new patch for version {version} in application "
-                f"package {self.package_name}? Once added, this operation cannot be undone."
-            )
-            if not policy.should_proceed(user_prompt):
-                if is_interactive:
-                    cc.message("Not creating a new patch.")
-                    raise typer.Exit(0)
-                else:
-                    cc.message(
-                        "Cannot create a new patch non-interactively without --force."
-                    )
-                    raise typer.Exit(1)
-
-        # Define a new version in the application package
-        if not self.get_existing_version_info(version):
-            self.add_new_version(version=version)
-            return  # A new version created automatically has patch 0, we do not need to further increment the patch.
-
-        # Add a new patch to an existing (old) version
-        self.add_new_patch_to_version(version=version, patch=patch)
 
 
 class NativeAppVersionDropProcessor(NativeAppManager, NativeAppCommandProcessor):
@@ -295,12 +100,7 @@ class NativeAppVersionDropProcessor(NativeAppManager, NativeAppCommandProcessor)
 
         # 1. Check for existing an existing application package
         show_obj_row = self.get_existing_app_pkg_info()
-        if show_obj_row:
-            # Check for the right owner role
-            ensure_correct_owner(
-                row=show_obj_row, role=self.package_role, obj_name=self.package_name
-            )
-        else:
+        if not show_obj_row:
             raise ApplicationPackageDoesNotExistError(self.package_name)
 
         # 2. Check distribution of the existing application package

snowflake/cli/_plugins/snowpark/commands.py

@@ -46,6 +46,10 @@ from snowflake.cli._plugins.snowpark.package.anaconda_packages import (
     AnacondaPackagesManager,
 )
 from snowflake.cli._plugins.snowpark.package.commands import app as package_app
+from snowflake.cli._plugins.snowpark.snowpark_entity_model import (
+    FunctionEntityModel,
+    ProcedureEntityModel,
+)
 from snowflake.cli._plugins.snowpark.snowpark_project_paths import (
     SnowparkProjectPaths,
 )
@@ -87,10 +91,6 @@ from snowflake.cli.api.output.types import (
 from snowflake.cli.api.project.definition_conversion import (
     convert_project_definition_to_v2,
 )
-from snowflake.cli.api.project.schemas.entities.snowpark_entity import (
-    FunctionEntityModel,
-    ProcedureEntityModel,
-)
 from snowflake.cli.api.project.schemas.project_definition import (
     ProjectDefinition,
     ProjectDefinitionV2,
@@ -446,5 +446,5 @@ def describe(
 def _get_v2_project_definition(cli_context) -> ProjectDefinitionV2:
     pd = cli_context.project_definition
     if not pd.meets_version_requirement("2"):
-        pd = convert_project_definition_to_v2(pd)
+        pd = convert_project_definition_to_v2(cli_context.project_root, pd)
     return pd

snowflake/cli/_plugins/snowpark/common.py

@@ -21,6 +21,10 @@ from typing import Dict, List, Set
 
 from click import UsageError
 from snowflake.cli._plugins.snowpark.models import Requirement
+from snowflake.cli._plugins.snowpark.snowpark_entity_model import (
+    ProcedureEntityModel,
+    SnowparkEntityModel,
+)
 from snowflake.cli._plugins.snowpark.snowpark_project_paths import Artefact
 from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.constants import (
@@ -30,10 +34,6 @@ from snowflake.cli.api.constants import (
     PROJECT_TEMPLATE_VARIABLE_OPENING,
     ObjectType,
 )
-from snowflake.cli.api.project.schemas.entities.snowpark_entity import (
-    ProcedureEntityModel,
-    SnowparkEntityModel,
-)
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
 

snowflake/cli/_plugins/snowpark/models.py

@@ -121,13 +121,14 @@ class WheelMetadata:
                 if line.startswith(dep_keyword)
             ]
             name = cls._get_name_from_wheel_filename(wheel_path.name)
+
             return cls(name=name, wheel_path=wheel_path, dependencies=dependencies)
 
     @staticmethod
     def _get_name_from_wheel_filename(wheel_filename: str) -> str:
         # wheel filename is in format {name}-{version}[-{extra info}]
         # https://peps.python.org/pep-0491/#file-name-convention
-        return wheel_filename.split("-")[0]
+        return wheel_filename.split("-")[0].lower()
 
     @staticmethod
     def to_wheel_name_format(package_name: str) -> str:

snowflake/cli/{api/entities → _plugins/snowpark}/snowpark_entity.py

@@ -1,10 +1,10 @@
 from typing import Generic, TypeVar
 
-from snowflake.cli.api.entities.common import EntityBase
-from snowflake.cli.api.project.schemas.entities.snowpark_entity import (
+from snowflake.cli._plugins.snowpark.snowpark_entity_model import (
     FunctionEntityModel,
     ProcedureEntityModel,
 )
+from snowflake.cli.api.entities.common import EntityBase
 
 T = TypeVar("T")
 

snowflake/cli/{api/project/schemas/entities/snowpark_entity.py → _plugins/snowpark/snowpark_entity_model.py}

@@ -22,12 +22,13 @@ from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.project.schemas.entities.common import (
     EntityModelBase,
     ExternalAccessBaseModel,
+    ImportsBaseModel,
 )
-from snowflake.cli.api.project.schemas.snowpark.argument import Argument
 from snowflake.cli.api.project.schemas.updatable_model import (
     DiscriminatorField,
     UpdatableModel,
 )
+from snowflake.cli.api.project.schemas.v1.snowpark.argument import Argument
 
 
 class PathMapping(UpdatableModel):
@@ -43,7 +44,7 @@ class PathMapping(UpdatableModel):
     )
 
 
-class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel):
+class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel, ImportsBaseModel):
     handler: str = Field(
         title="Function’s or procedure’s implementation of the object inside source module",
         examples=["functions.hello_function"],
@@ -57,10 +58,6 @@ class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel):
     runtime: Optional[Union[str, float]] = Field(
         title="Python version to use when executing ", default=None
     )
-    imports: Optional[List[str]] = Field(
-        title="Stage and path to previously uploaded files you want to import",
-        default=[],
-    )
     stage: str = Field(title="Stage in which artifacts will be stored")
     artifacts: List[Union[PathMapping, str]] = Field(title="List of required sources")
 

snowflake/cli/_plugins/snowpark/snowpark_project_paths.py

@@ -16,11 +16,11 @@ from __future__ import annotations
 from dataclasses import dataclass
 from pathlib import Path, PurePosixPath
 
+from snowflake.cli._plugins.snowpark.snowpark_entity_model import PathMapping
 from snowflake.cli._plugins.snowpark.zipper import zip_dir
 from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.constants import DEPLOYMENT_STAGE
 from snowflake.cli.api.identifiers import FQN
-from snowflake.cli.api.project.schemas.entities.snowpark_entity import PathMapping
 from snowflake.cli.api.secure_path import SecurePath
 
 

snowflake/cli/_plugins/stage/manager.py

@@ -57,6 +57,9 @@ EXECUTE_SUPPORTED_FILES_FORMATS = (
     ".py",
 )  # tuple to preserve order but it's a set
 
+# Replace magic numbers with constants
+OMIT_FIRST = slice(1, None)
+
 
 @dataclass
 class StagePathParts:
@@ -67,7 +70,7 @@ class StagePathParts:
 
     @classmethod
     def get_directory(cls, stage_path: str) -> str:
-        return "/".join(Path(stage_path).parts[1:])
+        return "/".join(Path(stage_path).parts[OMIT_FIRST])
 
     @property
     def path(self) -> str:
@@ -119,7 +122,9 @@ class DefaultStagePathParts(StagePathParts):
         self.directory = self.get_directory(stage_path)
         self.stage = StageManager.get_stage_from_path(stage_path)
         stage_name = self.stage.split(".")[-1]
-        stage_name = stage_name[1:] if stage_name.startswith("@") else stage_name
+        stage_name = (
+            stage_name[OMIT_FIRST] if stage_name.startswith("@") else stage_name
+        )
         self.stage_name = stage_name
         self.is_directory = True if stage_path.endswith("/") else False
 
@@ -133,7 +138,7 @@ class DefaultStagePathParts(StagePathParts):
 
     def replace_stage_prefix(self, file_path: str) -> str:
         stage = Path(self.stage).parts[0]
-        file_path_without_prefix = Path(file_path).parts[1:]
+        file_path_without_prefix = Path(file_path).parts[OMIT_FIRST]
         return f"{stage}/{'/'.join(file_path_without_prefix)}"
 
     def add_stage_prefix(self, file_path: str) -> str:
@@ -461,7 +466,7 @@ class StageManager(SqlExecutionMixin):
         on_error: OnErrorType,
     ) -> Dict:
         try:
-            query = f"execute immediate from {file_stage_path}"
+            query = f"execute immediate from {self.quote_stage_name(file_stage_path)}"
             if variables:
                 query += variables
             self._execute_query(query)

snowflake/cli/_plugins/streamlit/commands.py

@@ -26,6 +26,9 @@ from snowflake.cli._plugins.object.command_aliases import (
     scope_option,
 )
 from snowflake.cli._plugins.streamlit.manager import StreamlitManager
+from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
+    StreamlitEntityModel,
+)
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.decorators import (
     with_experimental_behaviour,
@@ -49,9 +52,6 @@ from snowflake.cli.api.output.types import (
 from snowflake.cli.api.project.definition_conversion import (
     convert_project_definition_to_v2,
 )
-from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
-    StreamlitEntityModel,
-)
 
 app = SnowTyperFactory(
     name="streamlit",
@@ -138,7 +138,7 @@ def streamlit_deploy(
             raise NoProjectDefinitionError(
                 project_type="streamlit", project_root=cli_context.project_root
             )
-        pd = convert_project_definition_to_v2(pd)
+        pd = convert_project_definition_to_v2(cli_context.project_root, pd)
 
     streamlits: Dict[str, StreamlitEntityModel] = pd.get_entities_by_type(
         entity_type="streamlit"

snowflake/cli/_plugins/streamlit/manager.py

@@ -18,21 +18,23 @@ import logging
 from pathlib import Path
 from typing import List, Optional
 
+from click import ClickException
 from snowflake.cli._plugins.connection.util import (
     MissingConnectionAccountError,
     MissingConnectionRegionError,
     make_snowsight_url,
 )
+from snowflake.cli._plugins.object.manager import ObjectManager
 from snowflake.cli._plugins.stage.manager import StageManager
+from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
+    StreamlitEntityModel,
+)
 from snowflake.cli.api.commands.experimental_behaviour import (
     experimental_behaviour_enabled,
 )
 from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.identifiers import FQN
-from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
-    StreamlitEntityModel,
-)
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
@@ -96,12 +98,16 @@ class StreamlitManager(SqlExecutionMixin):
             query.append(f"ROOT_LOCATION = '{from_stage_name}'")
 
         query.append(f"MAIN_FILE = '{streamlit.main_file}'")
-
+        if streamlit.imports:
+            query.append(streamlit.get_imports_sql())
         if streamlit.query_warehouse:
             query.append(f"QUERY_WAREHOUSE = {streamlit.query_warehouse}")
         if streamlit.title:
             query.append(f"TITLE = '{streamlit.title}'")
 
+        if streamlit.comment:
+            query.append(f"COMMENT = '{streamlit.comment}'")
+
         if streamlit.external_access_integrations:
             query.append(streamlit.get_external_access_integrations_sql())
 
@@ -112,6 +118,13 @@ class StreamlitManager(SqlExecutionMixin):
 
     def deploy(self, streamlit: StreamlitEntityModel, replace: bool = False):
         streamlit_id = streamlit.fqn.using_connection(self._conn)
+        if (
+            ObjectManager().object_exists(object_type="streamlit", fqn=streamlit_id)
+            and not replace
+        ):
+            raise ClickException(
+                f"Streamlit {streamlit.fqn} already exist. If you want to replace it use --replace flag."
+            )
 
         # for backwards compatibility - quoted stage path might be case-sensitive
         # https://docs.snowflake.com/en/sql-reference/identifiers-syntax#double-quoted-identifiers

snowflake/cli/{api/entities → _plugins/streamlit}/streamlit_entity.py

@@ -1,7 +1,7 @@
-from snowflake.cli.api.entities.common import EntityBase
-from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
+from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
     StreamlitEntityModel,
 )
+from snowflake.cli.api.entities.common import EntityBase
 
 
 class StreamlitEntity(EntityBase[StreamlitEntityModel]):

snowflake/cli/{api/project/schemas/entities → _plugins/streamlit}/streamlit_entity_model.py

@@ -20,17 +20,19 @@ from pydantic import Field, model_validator
 from snowflake.cli.api.project.schemas.entities.common import (
     EntityModelBase,
     ExternalAccessBaseModel,
+    ImportsBaseModel,
 )
 from snowflake.cli.api.project.schemas.updatable_model import (
     DiscriminatorField,
 )
 
 
-class StreamlitEntityModel(EntityModelBase, ExternalAccessBaseModel):
+class StreamlitEntityModel(EntityModelBase, ExternalAccessBaseModel, ImportsBaseModel):
     type: Literal["streamlit"] = DiscriminatorField()  # noqa: A003
     title: Optional[str] = Field(
         title="Human-readable title for the Streamlit dashboard", default=None
     )
+    comment: Optional[str] = Field(title="Comment for the Streamlit app", default=None)
     query_warehouse: str = Field(
         title="Snowflake warehouse to host the app", default=None
     )
@@ -48,23 +50,14 @@ class StreamlitEntityModel(EntityModelBase, ExternalAccessBaseModel):
         default=None,
     )
 
-    @model_validator(mode="after")
-    def main_file_must_be_in_artifacts(self):
-        if not self.artifacts:
-            return self
-
-        if Path(self.main_file) not in self.artifacts:
-            raise ValueError(
-                f"Specified main file {self.main_file} is not included in artifacts."
-            )
-        return self
-
     @model_validator(mode="after")
     def artifacts_must_exists(self):
         if not self.artifacts:
             return self
 
         for artifact in self.artifacts:
+            if "*" in artifact.name:
+                continue
             if not artifact.exists():
                 raise ValueError(
                     f"Specified artifact {artifact} does not exist locally."

snowflake/cli/_plugins/workspace/action_context.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional
+from typing import Callable, Optional
 
 from snowflake.cli.api.console.abc import AbstractConsole
 
@@ -15,3 +15,4 @@ class ActionContext:
     project_root: Path
     default_role: str
     default_warehouse: Optional[str]
+    get_entity: Callable