Extract find_entity function

src/snowflake/cli/_plugins/nativeapp/commands.py

@@ -22,7 +22,6 @@
 from typing import Generator, Iterable, List, Optional, cast
 
 import typer
-from click import ClickException
 from snowflake.cli._plugins.nativeapp.application_entity_model import (
     ApplicationEntityModel,
 )
@@ -53,6 +52,7 @@
     shallow_git_clone,
 )
 from snowflake.cli._plugins.nativeapp.v2_conversions.v2_to_v1_decorator import (
+    find_entity,
     nativeapp_definition_v2_to_v1,
 )
 from snowflake.cli._plugins.nativeapp.version.commands import app as versions_app
@@ -341,31 +341,18 @@ def app_teardown(
         processor.process(interactive, force, cascade)
     else:
         # New behaviour, multi-app aware so teardown all the apps created from the package
-        # Determine the package entity to drop, there must be one
-        app_package_definition: Optional[ApplicationPackageEntityModel] = None
-        packages: dict[
-            str, ApplicationPackageEntityModel
-        ] = project.get_entities_by_type(ApplicationPackageEntityModel.get_type())
-        if package_entity_id:
-            # If the user specified a package entity ID (or we inferred one from the app entity), use that one directly
-            app_package_definition = packages.get(package_entity_id)
-        elif len(packages) == 1:
-            # Otherwise, if there is only one package entity, fall back to that one
-            app_package_definition = next(iter(packages.values()))
-        elif len(packages) > 1:
-            # If there are multiple package entities, the user must specify which one to use
-            raise ClickException(
-                "More than one application package entity exists in the project definition file, "
-                "specify --package-entity-id to choose which one to operate on."
-            )
 
-        # If we don't have a package entity to convert, error out since it's not optional
-        if not app_package_definition:
-            with_id = f'with ID "{package_entity_id}" ' if package_entity_id else ""
-            raise ClickException(
-                f"Could not find an application package entity {with_id}in the project definition file."
-            )
+        # Determine the package entity to drop, there must be one
+        app_package_definition = find_entity(
+            project,
+            ApplicationPackageEntityModel,
+            package_entity_id,
+            disambiguation_option="--package-entity-id",
+            required=True,
+        )
+        assert app_package_definition is not None, "package entity is required"
 
+        # Same implementation as `snow ws drop`
         ws = WorkspaceManager(
             project_definition=cli_context.project_definition,
             project_root=cli_context.project_root,

src/snowflake/cli/_plugins/nativeapp/v2_conversions/v2_to_v1_decorator.py

@@ -16,7 +16,7 @@
 
 import inspect
 from functools import wraps
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Type, TypeVar, Union
 
 import typer
 from click import ClickException
@@ -31,6 +31,7 @@
     get_cli_context_manager,
 )
 from snowflake.cli.api.commands.decorators import _options_decorator_factory
+from snowflake.cli.api.project.schemas.entities.common import EntityModelBase
 from snowflake.cli.api.project.schemas.project_definition import (
     DefinitionV11,
     DefinitionV20,
@@ -59,31 +60,14 @@ def _pdf_v2_to_v1(
 ) -> DefinitionV11:
     pdfv1: Dict[str, Any] = {"definition_version": "1.1", "native_app": {}}
 
-    app_package_definition: Optional[ApplicationPackageEntityModel] = None
-    app_definition: Optional[ApplicationEntityModel] = None
-
-    # Enumerate all application package and application entities in the project definition
-    packages: dict[
-        str, ApplicationPackageEntityModel
-    ] = v2_definition.get_entities_by_type(ApplicationPackageEntityModel.get_type())
-    apps: dict[str, ApplicationEntityModel] = v2_definition.get_entities_by_type(
-        ApplicationEntityModel.get_type()
-    )
-
     # Determine the application entity to convert, there can be zero or one
-    if app_entity_id:
-        # If the user specified an app entity ID, use that one directly
-        app_definition = apps.get(app_entity_id)
-    elif len(apps) == 1:
-        # Otherwise, if there is only one app entity, fall back to that one
-        app_definition = next(iter(apps.values()))
-    elif len(apps) > 1 and app_required:
-        # If there are multiple app entities and the command being called requires one,
-        # the user must specify which one to use
-        raise ClickException(
-            "More than one application entity exists in the project definition file, "
-            "specify --app-entity-id to choose which one to operate on."
-        )
+    app_definition = find_entity(
+        v2_definition,
+        ApplicationEntityModel,
+        app_entity_id,
+        disambiguation_option="--app-entity-id",
+        required=app_required,
+    )
 
     # Infer or verify the package if we have an app entity to convert
     if app_definition:
@@ -99,36 +83,21 @@ def _pdf_v2_to_v1(
                     f"or omit the --package-entity-id flag to automatically use the package entity "
                     f"that the application entity targets."
                 )
-        elif target_package in packages:
+        elif target_package in v2_definition.get_entities_by_type(
+            ApplicationPackageEntityModel.get_type()
+        ):
             # If the user didn't target a specific package entity, use the one the app entity targets
             package_entity_id = target_package
-    elif app_required:
-        # If an app entity is required but we don't have one, error out
-        with_id = f'with ID "{app_entity_id}" ' if app_entity_id else ""
-        raise ClickException(
-            f"Could not find an application entity {with_id}in the project definition file."
-        )
 
     # Determine the package entity to convert, there must be one
-    if package_entity_id:
-        # If the user specified a package entity ID (or we inferred one from the app entity), use that one directly
-        app_package_definition = packages.get(package_entity_id)
-    elif len(packages) == 1:
-        # Otherwise, if there is only one package entity, fall back to that one
-        app_package_definition = next(iter(packages.values()))
-    elif len(packages) > 1:
-        # If there are multiple package entities, the user must specify which one to use
-        raise ClickException(
-            "More than one application package entity exists in the project definition file, "
-            "specify --package-entity-id to choose which one to operate on."
-        )
-
-    # If we don't have a package entity to convert, error out since it's not optional
-    if not app_package_definition:
-        with_id = f'with ID "{package_entity_id}" ' if package_entity_id else ""
-        raise ClickException(
-            f"Could not find an application package entity {with_id}in the project definition file."
-        )
+    app_package_definition = find_entity(
+        v2_definition,
+        ApplicationPackageEntityModel,
+        package_entity_id,
+        disambiguation_option="--package-entity-id",
+        required=True,
+    )
+    assert app_package_definition is not None, "package entity is required"
 
     # NativeApp
     if app_definition and app_definition.fqn.identifier:
@@ -188,6 +157,45 @@ def _pdf_v2_to_v1(
     return result.project_definition
 
 
+T = TypeVar("T", bound=EntityModelBase)
+
+
+def find_entity(
+    project_definition: DefinitionV20,
+    entity_class: Type[T],
+    entity_id: str,
+    disambiguation_option: str,
+    required: bool,
+) -> T | None:
+    entity_type = entity_class.get_type()
+    entities = project_definition.get_entities_by_type(entity_type)
+
+    entity: Optional[T] = None
+
+    # Determine the package entity to convert, there must be one
+    if entity_id:
+        # If the user specified a package entity ID (or we inferred one from the app entity), use that one directly
+        entity = entities.get(entity_id)
+    elif len(entities) == 1:
+        # Otherwise, if there is only one package entity, fall back to that one
+        entity = next(iter(entities.values()))
+    elif len(entities) > 1:
+        # If there are multiple package entities, the user must specify which one to use
+        raise ClickException(
+            f"More than one {entity_type} entity exists in the project definition file, "
+            f"specify {disambiguation_option} to choose which one to operate on."
+        )
+
+    # If we don't have a package entity to convert, error out since it's not optional
+    if not entity and required:
+        with_id = f'with ID "{entity_id}" ' if entity_id else ""
+        raise ClickException(
+            f"Could not find an {entity_type} entity {with_id}in the project definition file."
+        )
+
+    return entity
+
+
 def nativeapp_definition_v2_to_v1(*, app_required: bool = False):
     """
     A command decorator that attempts to automatically convert a native app project from

tests/__snapshots__/test_help_messages.ambr

@@ -789,11 +789,6 @@
   |                                                    operate when              |
   |                                                    definition_version is 2   |
   |                                                    or higher.                |
-  | --app-entity-id                              TEXT  The ID of the application |
-  |                                                    entity on which to        |
-  |                                                    operate when              |
-  |                                                    definition_version is 2   |
-  |                                                    or higher.                |
   | --project            -p                      TEXT  Path where Snowflake      |
   |                                                    project resides. Defaults |
   |                                                    to current working        |