Fix snowpark lookup responses

src/snowflake/cli/plugins/snowpark/commands.py

@@ -40,7 +40,7 @@
     check_if_replace_is_required,
 )
 from snowflake.cli.plugins.snowpark.manager import FunctionManager, ProcedureManager
-from snowflake.cli.plugins.snowpark.models import PypiOption
+from snowflake.cli.plugins.snowpark.models import YesNoAsk
 from snowflake.cli.plugins.snowpark.package_utils import get_snowflake_packages
 from snowflake.cli.plugins.snowpark.snowpark_shared import (
     CheckAnacondaForPyPiDependencies,
@@ -307,9 +307,9 @@ def _get_snowpark_artifact_path(snowpark_definition: Snowpark):
 @app.command("build")
 @with_project_definition("snowpark")
 def build(
-    pypi_download: PypiOption = PyPiDownloadOption,
+    pypi_download: YesNoAsk = PyPiDownloadOption,
     check_anaconda_for_pypi_deps: bool = CheckAnacondaForPyPiDependencies,
-    package_native_libraries: PypiOption = PackageNativeLibrariesOption,
+    package_native_libraries: YesNoAsk = PackageNativeLibrariesOption,
     **options,
 ) -> CommandResult:
     """

src/snowflake/cli/plugins/snowpark/models.py

@@ -7,7 +7,7 @@
 from requirements.requirement import Requirement
 
 
-class PypiOption(Enum):
+class YesNoAsk(Enum):
     YES = "yes"
     NO = "no"
     ASK = "ask"

src/snowflake/cli/plugins/snowpark/package/commands.py

@@ -2,20 +2,28 @@
 
 import logging
 from pathlib import Path
+from textwrap import dedent
 
 import typer
 from snowflake.cli.api.commands.snow_typer import SnowTyper
+from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.output.types import CommandResult, MessageResult
+from snowflake.cli.plugins.snowpark import package_utils
+from snowflake.cli.plugins.snowpark.models import YesNoAsk
 from snowflake.cli.plugins.snowpark.package.manager import (
+    check_if_package_in_anaconda,
     cleanup_after_install,
-    create,
+    create_package,
     lookup,
     upload,
 )
 from snowflake.cli.plugins.snowpark.package.utils import (
-    CreatedSuccessfully,
-    NotInAnaconda,
-    RequiresPackages,
+    get_readable_list_of_requirements,
+)
+from snowflake.cli.plugins.snowpark.package_utils import check_for_native_libraries
+from snowflake.cli.plugins.snowpark.snowpark_shared import (
+    PackageNativeLibrariesOption,
+    check_if_can_continue_with_native_libs,
 )
 
 app = SnowTyper(
@@ -44,18 +52,18 @@
 def package_lookup(
     name: str = typer.Argument(..., help="Name of the package."),
     install_packages: bool = install_option,
-    deprecated_install_option: bool = deprecated_install_option,
+    _deprecated_install_option: bool = deprecated_install_option,
     **options,
 ) -> CommandResult:
     """
     Checks if a package is available on the Snowflake Anaconda channel.
     If the `--pypi-download` flag is provided, this command checks all dependencies of the packages
     outside Snowflake channel.
     """
-    if deprecated_install_option:
-        install_packages = deprecated_install_option
+    if _deprecated_install_option:
+        install_packages = _deprecated_install_option
 
-    lookup_result = lookup(name=name, install_packages=install_packages)
+    lookup_result = lookup(package_name=name, install_packages=install_packages)
     return MessageResult(lookup_result.message)
 
 
@@ -91,31 +99,65 @@ def package_upload(
 @app.command("create", requires_connection=True)
 @cleanup_after_install
 def package_create(
-    name: str = typer.Argument(
+    package_name: str = typer.Argument(
         ...,
         help="Name of the package to create.",
     ),
     install_packages: bool = install_option,
-    deprecated_install_option: bool = deprecated_install_option,
+    package_native_libraries: YesNoAsk = PackageNativeLibrariesOption,
+    _deprecated_install_option: bool = deprecated_install_option,
     **options,
 ) -> CommandResult:
     """
     Creates a Python package as a zip file that can be uploaded to a stage and imported for a Snowpark Python app.
     """
-    if deprecated_install_option:
-        install_packages = deprecated_install_option
-
-    if (
-        type(lookup_result := lookup(name=name, install_packages=install_packages))
-        in [
-            NotInAnaconda,
-            RequiresPackages,
-        ]
-        and type(creation_result := create(name)) == CreatedSuccessfully
-    ):
-        message = creation_result.message
-        if type(lookup_result) == RequiresPackages:
-            message += "\n" + lookup_result.message
-    else:
-        message = lookup_result.message
+    if _deprecated_install_option:
+        install_packages = _deprecated_install_option
+
+    with cli_console.phase("Anaconda check"):
+        available_in_anaconda, _ = check_if_package_in_anaconda(package_name)
+
+        # If the package is in Anaconda there's nothing to do
+        if available_in_anaconda:
+            return MessageResult(
+                f"Package {package_name} is available on the Snowflake Anaconda channel."
+            )
+        cli_console.message("Package not available on Anaconda")
+
+    # Now there are some missing package in Anaconda and if we are not
+    # allowed to install them, so we can only quit
+    if not install_packages:
+        return MessageResult(
+            f"Package {package_name} is not available on Anaconda. Creating the package "
+            f"requires installing additional packages. Enable it by providing --pypi-download."
+        )
+
+    with cli_console.phase("Building the package"):
+        split_requirements = package_utils.install_packages(
+            perform_anaconda_check=True, package_name=package_name, file_name=None
+        )
+
+        if check_for_native_libraries():
+            check_if_can_continue_with_native_libs(package_native_libraries)
+
+        # If we can do a local installation, then lookup created packages directory that we now zip
+        zip_file = create_package(package_name)
+
+    message = dedent(
+        f"""
+    Package {package_name}.zip created. Upload it to stage using:
+    snow snowpark package upload -f {zip_file} -s <stage-name>
+    and reference it in your procedure or function imports.
+    """
+    )
+
+    # Some requirements can be available in Snowflake and not included in the zip.
+    # In such case users will have to add them manually, so we add the list to the message
+    if split_requirements.snowflake:
+        # There are requirements available in snowflake that has to be added to package list
+        packages_to_be_added_manually = get_readable_list_of_requirements(
+            split_requirements.snowflake
+        )
+        message += f"\nYou should also include following packages in your function or procedure:\n{packages_to_be_added_manually}"
+
     return MessageResult(message)

src/snowflake/cli/plugins/snowpark/package/manager.py

@@ -1,48 +1,58 @@
 from __future__ import annotations
 
 import logging
-import os.path
 from functools import wraps
 from pathlib import Path
 
 from requirements.requirement import Requirement
+from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.constants import PACKAGES_DIR
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.cli.plugins.object.stage.manager import StageManager
 from snowflake.cli.plugins.snowpark import package_utils
-from snowflake.cli.plugins.snowpark.models import SplitRequirements
 from snowflake.cli.plugins.snowpark.package.utils import (
-    CreatedSuccessfully,
     InAnaconda,
     LookupResult,
-    NothingFound,
     NotInAnaconda,
-    RequiresPackages,
+    NotInAnacondaButRequiresNativePackages,
     prepare_app_zip,
 )
+from snowflake.cli.plugins.snowpark.package_utils import check_for_native_libraries
 from snowflake.cli.plugins.snowpark.zipper import zip_dir
 
 log = logging.getLogger(__name__)
 
 
-def lookup(name: str, install_packages: bool) -> LookupResult:
+def check_if_package_in_anaconda(package_name):
+    cli_console.step("Checking package availability in Snowflake Anaconda")
+    split_requirements = package_utils.parse_anaconda_packages(
+        [Requirement.parse(package_name)]
+    )
+    available_in_anaconda = (
+        split_requirements.snowflake and not split_requirements.other
+    )
+    return available_in_anaconda, split_requirements
 
-    package_response = package_utils.parse_anaconda_packages([Requirement.parse(name)])
 
-    if package_response.snowflake and not package_response.other:
-        return InAnaconda(package_response, name)
-    elif install_packages:
-        status, result = package_utils.install_packages(
-            perform_anaconda_check=True, package_name=name, file_name=None
-        )
+def lookup(package_name: str, install_packages: bool) -> LookupResult:
+    available_in_anaconda, split_requirements = check_if_package_in_anaconda(
+        package_name
+    )
+    # If all packages are in snowflake then there's no need for user action
+    if available_in_anaconda:
+        return InAnaconda(requirements=split_requirements, name=package_name)
 
-        if status:
-            if result.snowflake:
-                return RequiresPackages(result, name)
-            else:
-                return NotInAnaconda(result, name)
+    # There are missing packages, but we are not allowed to install them
+    if split_requirements.other and not install_packages:
+        return NotInAnaconda(requirements=split_requirements, name=package_name)
 
-    return NothingFound(SplitRequirements([], []), name)
+    split_requirements = package_utils.install_packages(
+        perform_anaconda_check=True, package_name=package_name, file_name=None
+    )
+    if check_for_native_libraries():
+        return NotInAnacondaButRequiresNativePackages(split_requirements, package_name)
+
+    return NotInAnaconda(split_requirements, package_name)
 
 
 def upload(file: Path, stage: str, overwrite: bool):
@@ -56,20 +66,18 @@ def upload(file: Path, stage: str, overwrite: bool):
             temp_app_zip_path.path, stage, overwrite=overwrite
         ).fetchone()
 
-    message = f"Package {file} {put_response[6]} to Snowflake @{stage}/{file}."
-
     if put_response[6] == "SKIPPED":
-        message = "Package already exists on stage. Consider using --overwrite to overwrite the file."
+        return "Package already exists on stage. Consider using --overwrite to overwrite the file."
 
-    return message
+    return f"Package {file} {put_response[6]} to Snowflake. Add '@{stage}/{file}' to imports of your function or procedure."
 
 
-def create(zip_name: str):
+def create_package(zip_name: str):
     file_name = zip_name if zip_name.endswith(".zip") else f"{zip_name}.zip"
-    zip_dir(dest_zip=Path(file_name), source=Path.cwd() / ".packages")
-
-    if os.path.exists(file_name):
-        return CreatedSuccessfully(zip_name, Path(file_name))
+    file_path = Path(file_name)
+    cli_console.step(f"Creating {file_path}")
+    zip_dir(dest_zip=file_path, source=Path.cwd() / ".packages")
+    return file_path
 
 
 def cleanup_after_install(func):

src/snowflake/cli/plugins/snowpark/package/utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from pathlib import Path
+from textwrap import dedent
 from typing import List
 
 from requirements.requirement import Requirement
@@ -25,39 +25,27 @@ def message(self):
         return f"Package {self.name} is available on the Snowflake Anaconda channel."
 
 
-class RequiresPackages(LookupResult):
-    @property
-    def message(self):
-        return f"""The package {self.name} is supported, but does depend on the
-                following Snowflake supported libraries. You should
-                include the following in your packages:
-                {get_readable_list_of_requirements(self.requirements.snowflake)}"""
-
-
 class NotInAnaconda(LookupResult):
     @property
     def message(self):
-        return f"""The package {self.name} is avaiable through PIP. You can create a zip using:\n
-                snow snowpark package create {self.name} --pypi-download"""
-
-
-class NothingFound(LookupResult):
-    @property
-    def message(self):
-        return f"""Nothing found for {self.name}. Most probably, package is not avaiable on Snowflake Anaconda channel\n
-                    Please check the package name or try again with --pypi-download option"""
-
-
-@dataclass
-class CreateResult:
-    package_name: str
-    file_name: Path = Path()
+        return dedent(
+            f"""
+        The package {self.name} is not available on Snowflake Anaconda. You can create a zip using:
+        snow snowpark package create {self.name} --pypi-download
+        """
+        )
 
 
-class CreatedSuccessfully(CreateResult):
+class NotInAnacondaButRequiresNativePackages(NotInAnaconda):
     @property
     def message(self):
-        return f"Package {self.package_name}.zip created. You can now upload it to a stage (`snow snowpark package upload -f {self.package_name}.zip -s <stage-name>`) and reference it in your procedure or function."
+        return dedent(
+            f"""
+        The package {self.name} is not available on Snowflake Anaconda and requires native libraries.
+        You can try to create a zip using:
+        snow snowpark package create {self.name} --pypi-download
+        """
+        )
 
 
 def prepare_app_zip(file_path: SecurePath, temp_dir: SecurePath) -> SecurePath: