snowflake-cli 3.2.2__py3-none-any.whl → 3.4.1__py3-none-any.whl

snowflake/cli/_plugins/snowpark/snowpark_entity.py

@@ -1,16 +1,241 @@
-from typing import Generic, TypeVar
+from enum import Enum
+from pathlib import Path
+from typing import Generic, List, Optional, TypeVar
 
+from click import ClickException
+from snowflake.cli._plugins.nativeapp.feature_flags import FeatureFlag
+from snowflake.cli._plugins.snowpark import package_utils
+from snowflake.cli._plugins.snowpark.common import (
+    DEFAULT_RUNTIME,
+    map_path_mapping_to_artifact,
+    zip_and_copy_artifacts_to_deploy,
+)
+from snowflake.cli._plugins.snowpark.package.anaconda_packages import (
+    AnacondaPackages,
+    AnacondaPackagesManager,
+)
+from snowflake.cli._plugins.snowpark.package_utils import (
+    DownloadUnavailablePackagesResult,
+)
 from snowflake.cli._plugins.snowpark.snowpark_entity_model import (
     FunctionEntityModel,
     ProcedureEntityModel,
 )
+from snowflake.cli._plugins.snowpark.snowpark_project_paths import SnowparkProjectPaths
+from snowflake.cli._plugins.snowpark.zipper import zip_dir
+from snowflake.cli._plugins.workspace.context import ActionContext
 from snowflake.cli.api.entities.common import EntityBase
+from snowflake.cli.api.secure_path import SecurePath
+from snowflake.connector import ProgrammingError
 
 T = TypeVar("T")
 
 
+class CreateMode(
+    str, Enum
+):  # This should probably be moved to some common place, think where
+    create = "CREATE"
+    create_or_replace = "CREATE OR REPLACE"
+    create_if_not_exists = "CREATE IF NOT EXISTS"
+
+
 class SnowparkEntity(EntityBase[Generic[T]]):
-    pass
+    def __init__(self, *args, **kwargs):
+
+        if not FeatureFlag.ENABLE_NATIVE_APP_CHILDREN.is_enabled():
+            raise NotImplementedError("Snowpark entity is not implemented yet")
+        super().__init__(*args, **kwargs)
+
+    def action_bundle(
+        self,
+        action_ctx: ActionContext,
+        ignore_anaconda: bool,
+        skip_version_check: bool,
+        output_dir: Path | None = None,
+        index_url: str | None = None,
+        allow_shared_libraries: bool = False,
+        *args,
+        **kwargs,
+    ) -> List[Path]:
+        return self.bundle(
+            ignore_anaconda,
+            skip_version_check,
+            output_dir,
+            index_url,
+            allow_shared_libraries,
+        )
+
+    def action_deploy(
+        self,
+        action_ctx: ActionContext,
+        mode: CreateMode = CreateMode.create,
+        *args,
+        **kwargs,
+    ):
+        # TODO: After introducing bundle map, we should introduce file copying part here
+        return self.deploy(mode, *args, **kwargs)
+
+    def action_drop(self, action_ctx: ActionContext, *args, **kwargs):
+        return self._execute_query(self.get_drop_sql())
+
+    def action_describe(self, action_ctx: ActionContext, *args, **kwargs):
+        return self._execute_query(self.get_describe_sql())
+
+    def action_execute(
+        self,
+        action_ctx: ActionContext,
+        execution_arguments: List[str] | None = None,
+        *args,
+        **kwargs,
+    ):
+        return self._execute_query(self.get_execute_sql(execution_arguments))
+
+    def bundle(
+        self,
+        ignore_anaconda: bool,
+        skip_version_check: bool,
+        output_dir: Path | None = None,
+        index_url: str | None = None,
+        allow_shared_libraries: bool = False,
+    ) -> List[Path]:
+        """
+        Bundles the entity artifacts and dependencies into a directory.
+        Parameters:
+            output_dir: The directory to output the bundled artifacts to. Defaults to output dir in project root
+            ignore_anaconda: If True, ignores anaconda check and tries to download all packages using pip
+            skip_version_check: If True, skips version check when downloading packages
+            index_url: The index URL to use when downloading packages, if none set - default pip index is used (in most cases- Pypi)
+            allow_shared_libraries: If not set to True, using dependency with .so/.dll files will raise an exception
+        Returns:
+        """
+        # 0 Create a directory for the entity
+        project_paths = SnowparkProjectPaths(
+            project_root=self.root.absolute(),
+        )
+        if not output_dir:
+            output_dir = project_paths.bundle_root
+        if not output_dir.exists():  # type: ignore[union-attr]
+            SecurePath(output_dir).mkdir(parents=True)
+
+        # 1 Check if requirements exits
+        if (self.root / "requirements.txt").exists():
+            download_results = self._process_requirements(
+                bundle_dir=output_dir,  # type: ignore
+                archive_name="dependencies.zip",
+                requirements_file=SecurePath(self.root / "requirements.txt"),
+                ignore_anaconda=ignore_anaconda,
+                skip_version_check=skip_version_check,
+                index_url=index_url,
+                allow_shared_libraries=allow_shared_libraries,
+            )
+
+        # 2 get the artifacts list
+        artifacts = map_path_mapping_to_artifact(project_paths, self.model.artifacts)
+        from snowflake.cli.api.feature_flags import FeatureFlag
+
+        if FeatureFlag.ENABLE_SNOWPARK_GLOB_SUPPORT.is_enabled():
+            return zip_and_copy_artifacts_to_deploy(
+                artifacts, project_paths.bundle_root
+            )
+        else:
+            copied_files = []
+            for artifact in artifacts:
+                artifact.build()
+                copied_files.append(artifact.post_build_path)
+            return copied_files
+
+    def check_if_exists(
+        self, action_ctx: ActionContext
+    ) -> bool:  # TODO it should return current state, so we know if update is necessary
+        try:
+            current_state = self.action_describe(action_ctx)
+            return True
+        except ProgrammingError:
+            return False
+
+    def deploy(self, mode: CreateMode = CreateMode.create, *args, **kwargs):
+        return self._execute_query(self.get_deploy_sql(mode))
+
+    def get_deploy_sql(self, mode: CreateMode):
+        query = [
+            f"{mode.value} {self.model.type.upper()} {self.identifier}",
+            "COPY GRANTS",
+            f"RETURNS {self.model.returns}",
+            f"LANGUAGE PYTHON",
+            f"RUNTIME_VERSION '{self.model.runtime or DEFAULT_RUNTIME}'",
+            f"IMPORTS={','.join(self.model.imports)}",  # TODO: Add source files here after introducing bundlemap
+            f"HANDLER='{self.model.handler}'",
+        ]
+
+        if self.model.external_access_integrations:
+            query.append(self.model.get_external_access_integrations_sql())
+
+        if self.model.secrets:
+            query.append(self.model.get_secrets_sql())
+
+        if self.model.type == "procedure" and self.model.execute_as_caller:
+            query.append("EXECUTE AS CALLER")
+
+        return "\n".join(query)
+
+    def get_execute_sql(self, execution_arguments: List[str] | None = None):
+        raise NotImplementedError
+
+    def _process_requirements(  # TODO: maybe leave all the logic with requirements here - so download, write requirements file etc.
+        self,
+        bundle_dir: Path,
+        archive_name: str,  # TODO: not the best name, think of something else
+        requirements_file: Optional[SecurePath],
+        ignore_anaconda: bool,
+        skip_version_check: bool = False,
+        index_url: Optional[str] = None,
+        allow_shared_libraries: bool = False,
+    ) -> DownloadUnavailablePackagesResult:
+        """
+        Processes the requirements file and downloads the dependencies
+        Parameters:
+
+        """
+        anaconda_packages_manager = AnacondaPackagesManager()
+        with SecurePath.temporary_directory() as tmp_dir:
+            requirements = package_utils.parse_requirements(requirements_file)
+            anaconda_packages = (
+                AnacondaPackages.empty()
+                if ignore_anaconda
+                else anaconda_packages_manager.find_packages_available_in_snowflake_anaconda()
+            )
+            download_result = package_utils.download_unavailable_packages(
+                requirements=requirements,
+                target_dir=tmp_dir,
+                anaconda_packages=anaconda_packages,
+                skip_version_check=skip_version_check,
+                pip_index_url=index_url,
+            )
+
+            if download_result.anaconda_packages:
+                anaconda_packages.write_requirements_file_in_snowflake_format(
+                    file_path=SecurePath(bundle_dir / "requirements.txt"),
+                    requirements=download_result.anaconda_packages,
+                )
+
+            if download_result.downloaded_packages_details:
+                if (
+                    package_utils.detect_and_log_shared_libraries(
+                        download_result.downloaded_packages_details
+                    )
+                    and not allow_shared_libraries
+                ):
+                    raise ClickException(
+                        "Some packages contain shared (.so/.dll) libraries. "
+                        "Try again with allow_shared_libraries_flag."
+                    )
+
+                zip_dir(
+                    source=tmp_dir,
+                    dest_zip=bundle_dir / archive_name,
+                )
+
+        return download_result
 
 
 class FunctionEntity(SnowparkEntity[FunctionEntityModel]):
@@ -18,7 +243,17 @@ class FunctionEntity(SnowparkEntity[FunctionEntityModel]):
     A single UDF
     """
 
-    pass
+    # TO THINK OF
+    # Where will we get imports? Should we rely on bundle map? Or should it be self-sufficient in this matter?
+
+    def get_execute_sql(
+        self, execution_arguments: List[str] | None = None, *args, **kwargs
+    ):
+        if not execution_arguments:
+            execution_arguments = []
+        return (
+            f"SELECT {self.fqn}({', '.join([str(arg) for arg in execution_arguments])})"
+        )
 
 
 class ProcedureEntity(SnowparkEntity[ProcedureEntityModel]):
@@ -26,4 +261,12 @@ class ProcedureEntity(SnowparkEntity[ProcedureEntityModel]):
     A stored procedure
     """
 
-    pass
+    def get_execute_sql(
+        self,
+        execution_arguments: List[str] | None = None,
+    ):
+        if not execution_arguments:
+            execution_arguments = []
+        return (
+            f"CALL {self.fqn}({', '.join([str(arg) for arg in execution_arguments])})"
+        )

snowflake/cli/_plugins/snowpark/snowpark_entity_model.py

@@ -14,37 +14,27 @@
 
 from __future__ import annotations
 
-from pathlib import Path
+import glob
 from typing import List, Literal, Optional, Union
 
 from pydantic import Field, field_validator
+from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.project.schemas.entities.common import (
-    EntityModelBase,
+    EntityModelBaseWithArtifacts,
     ExternalAccessBaseModel,
     ImportsBaseModel,
+    PathMapping,
 )
 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):
-    class Config:
-        frozen = True
-
-    src: Path = Field(title="Source path (relative to project root)", default=None)
-
-    dest: Optional[str] = Field(
-        title="Destination path on stage",
-        description="Paths are relative to stage root; paths ending with a slash indicate that the destination is a directory which source files should be copied into.",
-        default=None,
-    )
-
-
-class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel, ImportsBaseModel):
+class SnowparkEntityModel(
+    EntityModelBaseWithArtifacts, ExternalAccessBaseModel, ImportsBaseModel
+):
     handler: str = Field(
         title="Function’s or procedure’s implementation of the object inside source module",
         examples=["functions.hello_function"],
@@ -59,17 +49,23 @@ class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel, ImportsBaseM
         title="Python version to use when executing ", default=None
     )
     stage: str = Field(title="Stage in which artifacts will be stored")
-    artifacts: List[Union[PathMapping, str]] = Field(title="List of required sources")
 
     @field_validator("artifacts")
     @classmethod
     def _convert_artifacts(cls, artifacts: Union[dict, str]):
         _artifacts = []
-        for artefact in artifacts:
-            if isinstance(artefact, PathMapping):
-                _artifacts.append(artefact)
+        for artifact in artifacts:
+            if (
+                (isinstance(artifact, str) and glob.has_magic(artifact))
+                or (isinstance(artifact, PathMapping) and glob.has_magic(artifact.src))
+            ) and FeatureFlag.ENABLE_SNOWPARK_GLOB_SUPPORT.is_disabled():
+                raise ValueError(
+                    "If you want to use glob patterns in artifacts, you need to enable the Snowpark new build feature flag (enable_snowpark_glob_support=true)"
+                )
+            if isinstance(artifact, PathMapping):
+                _artifacts.append(artifact)
             else:
-                _artifacts.append(PathMapping(src=artefact))
+                _artifacts.append(PathMapping(src=artifact))
         return _artifacts
 
     @field_validator("runtime")
@@ -79,14 +75,6 @@ class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel, ImportsBaseM
             return str(runtime_input)
         return runtime_input
 
-    @field_validator("artifacts")
-    @classmethod
-    def validate_artifacts(cls, artifacts: List[Path]) -> List[Path]:
-        for artefact in artifacts:
-            if "*" in str(artefact):
-                raise ValueError("Glob patterns not supported for Snowpark artifacts.")
-        return artifacts
-
     @property
     def udf_sproc_identifier(self) -> UdfSprocIdentifier:
         return UdfSprocIdentifier.from_definition(self)

snowflake/cli/_plugins/snowpark/snowpark_project_paths.py

@@ -13,38 +13,60 @@
 # limitations under the License.
 from __future__ import annotations
 
+import glob
+import os
+import re
 from dataclasses import dataclass
 from pathlib import Path, PurePosixPath
+from typing import Optional
 
-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.feature_flags import FeatureFlag
 from snowflake.cli.api.identifiers import FQN
+from snowflake.cli.api.project.project_paths import ProjectPaths, bundle_root
+from snowflake.cli.api.project.schemas.entities.common import PathMapping
 from snowflake.cli.api.secure_path import SecurePath
 
 
 @dataclass
-class SnowparkProjectPaths:
+class SnowparkProjectPaths(ProjectPaths):
     """
-    This class represents allows you to manage files paths related to given project.
+    This class allows you to manage files paths related to given project.
     """
 
-    project_root: Path
-
     def path_relative_to_root(self, artifact_path: Path) -> Path:
         if artifact_path.is_absolute():
             return artifact_path
         return (self.project_root / artifact_path).resolve()
 
-    def get_artefact_dto(self, artifact_path: PathMapping) -> Artefact:
-        return Artefact(
-            dest=artifact_path.dest,
-            path=self.path_relative_to_root(artifact_path.src),
-        )
-
-    def get_dependencies_artefact(self) -> Artefact:
-        return Artefact(dest=None, path=self.dependencies)
+    def get_artifact_dto(self, artifact_path: PathMapping) -> Artifact:
+        if FeatureFlag.ENABLE_SNOWPARK_GLOB_SUPPORT.is_enabled():
+            return Artifact(
+                project_root=self.project_root,
+                bundle_root=self.bundle_root,
+                dest=artifact_path.dest,
+                path=Path(artifact_path.src),
+            )
+        else:
+            return ArtifactOldBuild(
+                dest=artifact_path.dest,
+                path=self.path_relative_to_root(Path(artifact_path.src)),
+            )
+
+    def get_dependencies_artifact(self) -> Artifact:
+        if FeatureFlag.ENABLE_SNOWPARK_GLOB_SUPPORT.is_enabled():
+            return Artifact(
+                project_root=self.project_root,
+                bundle_root=self.bundle_root,
+                dest=None,
+                path=Path("dependencies.zip"),
+            )
+        else:
+            return ArtifactOldBuild(
+                dest=None, path=self.path_relative_to_root(Path("dependencies.zip"))
+            )
 
     @property
     def snowflake_requirements(self) -> SecurePath:
@@ -57,19 +79,130 @@ class SnowparkProjectPaths:
         return SecurePath(self.path_relative_to_root(Path("requirements.txt")))
 
     @property
-    def dependencies(self) -> Path:
-        return self.path_relative_to_root(Path("dependencies.zip"))
+    def bundle_root(self) -> Path:
+        return bundle_root(self.project_root, "snowpark")
 
 
 @dataclass(unsafe_hash=True)
-class Artefact:
-    """Helper for getting paths related to given artefact."""
+class Artifact:
+    """Helper for getting paths related to given artifact."""
 
+    project_root: Path
+    bundle_root: Path
     path: Path
     dest: str | None = None
 
+    def __init__(
+        self,
+        project_root: Path,
+        bundle_root: Path,
+        path: Path,
+        dest: Optional[str] = None,
+    ) -> None:
+        self.project_root = project_root
+        self.bundle_root = bundle_root
+        self.path = path
+        self.dest = dest
+        if self.dest and not self._is_dest_a_file() and not self.dest.endswith("/"):
+            self.dest = self.dest + "/"
+
+    @property
+    def _artifact_name(self) -> str:
+        """
+        Returns artifact name. Directories are mapped to corresponding .zip files.
+        For paths with glob patterns, the last part of the path is used.
+        For files, the file name is used.
+        """
+        if glob.has_magic(str(self.path)):
+            last_part = None
+            for part in self.path.parts:
+                if glob.has_magic(part):
+                    break
+                else:
+                    last_part = part
+            if not last_part:
+                last_part = os.path.commonpath(
+                    [str(self.path), str(self.path.absolute())]
+                )
+            return last_part + ".zip"
+        if (self.project_root / self.path).is_dir():
+            return self.path.stem + ".zip"
+        if (self.project_root / self.path).is_file() and self._is_dest_a_file():
+            return Path(self.dest).name  # type: ignore
+        return self.path.name
+
+    @property
+    def post_build_path(self) -> Path:
+        """
+        Returns post-build artifact path. Directories are mapped to corresponding .zip files.
+        """
+        bundle_root = self.bundle_root
+        path = (
+            self._path_until_asterisk()
+            if glob.has_magic(str(self.path))
+            else self.path.parent
+        )
+        if self._is_dest_a_file():
+            return bundle_root / self.dest  # type: ignore
+        return bundle_root / (self.dest or path) / self._artifact_name
+
+    def upload_path(self, stage: FQN | str | None) -> str:
+        """
+        Path on stage to which the artifact should be uploaded.
+        """
+        stage = stage or DEPLOYMENT_STAGE
+        if isinstance(stage, str):
+            stage = FQN.from_stage(stage).using_context()
+
+        stage_path = PurePosixPath(f"@{stage}")
+        if self.dest:
+            stage_path /= (
+                PurePosixPath(self.dest).parent if self._is_dest_a_file() else self.dest
+            )
+        else:
+            stage_path /= (
+                self._path_until_asterisk()
+                if glob.has_magic(str(self.path))
+                else PurePosixPath(self.path).parent
+            )
+
+        return str(stage_path) + "/"
+
+    def import_path(self, stage: FQN | str | None) -> str:
+        """Path for UDF/sproc imports clause."""
+        return self.upload_path(stage) + self._artifact_name
+
+    def _is_dest_a_file(self) -> bool:
+        if not self.dest:
+            return False
+        return re.search(r"\.[a-zA-Z0-9]{2,4}$", self.dest) is not None
+
+    def _path_until_asterisk(self) -> Path:
+        path = []
+        for part in self.path.parts:
+            if glob.has_magic(part):
+                break
+            else:
+                path.append(part)
+        return Path(*path[:-1])
+
+    # Can be removed after removing ENABLE_SNOWPARK_GLOB_SUPPORT feature flag.
+    def build(self) -> None:
+        raise NotImplementedError("Not implemented in Artifact class.")
+
+
+@dataclass(unsafe_hash=True)
+class ArtifactOldBuild(Artifact):
+    """Helper for getting paths related to given artifact."""
+
+    path: Path
+    dest: str | None = None
+
+    def __init__(self, path: Path, dest: Optional[str] = None) -> None:
+        super().__init__(project_root=Path(), bundle_root=Path(), path=path, dest=dest)
+
     @property
-    def _artefact_name(self) -> str:
+    def _artifact_name(self) -> str:
         if self.path.is_dir():
             return self.path.stem + ".zip"
         return self.path.name
@@ -77,13 +210,13 @@ class Artefact:
     @property
     def post_build_path(self) -> Path:
         """
-        Returns post-build artefact path. Directories are mapped to corresponding .zip files.
+        Returns post-build artifact path. Directories are mapped to corresponding .zip files.
         """
-        return self.path.parent / self._artefact_name
+        return self.path.parent / self._artifact_name
 
     def upload_path(self, stage: FQN | str | None) -> str:
         """
-        Path on stage to which the artefact should be uploaded.
+        Path on stage to which the artifact should be uploaded.
         """
         stage = stage or DEPLOYMENT_STAGE
         if isinstance(stage, str):
@@ -96,10 +229,10 @@ class Artefact:
 
     def import_path(self, stage: FQN | str | None) -> str:
         """Path for UDF/sproc imports clause."""
-        return self.upload_path(stage) + self._artefact_name
+        return self.upload_path(stage) + self._artifact_name
 
     def build(self) -> None:
-        """Build the artefact. Applies only to directories. Files are untouched."""
+        """Build the artifact. Applies only to directories. Files are untouched."""
         if not self.path.is_dir():
             return
         cli_console.step(f"Creating: {self.post_build_path.name}")

snowflake/cli/_plugins/snowpark/zipper.py

@@ -20,6 +20,10 @@ from pathlib import Path
 from typing import Dict, List, Literal
 from zipfile import ZIP_DEFLATED, ZipFile
 
+from snowflake.cli.api.artifacts.bundle_map import BundleMap
+from snowflake.cli.api.console import cli_console
+from snowflake.cli.api.secure_path import SecurePath
+
 log = logging.getLogger(__name__)
 
 IGNORED_FILES = [
@@ -64,7 +68,10 @@ def zip_dir(
     mode: Literal["r", "w", "x", "a"] = "w",
 ) -> None:
 
-    if isinstance(source, Path):
+    if not dest_zip.parent.exists():
+        SecurePath(dest_zip).parent.mkdir(parents=True)
+
+    if isinstance(source, Path) or isinstance(source, SecurePath):
         source = [source]
 
     files_to_pack: Dict[Path, List[Path]] = {
@@ -74,11 +81,36 @@ def zip_dir(
 
     with ZipFile(dest_zip, mode, ZIP_DEFLATED, allowZip64=True) as package_zip:
         for src, files in files_to_pack.items():
+            if isinstance(src, SecurePath):
+                src = src.path
             for file in files:
                 log.debug("Adding %s to %s", file, dest_zip)
                 package_zip.write(file, arcname=file.relative_to(src))
 
 
+def zip_dir_using_bundle_map(
+    bundle_map: BundleMap,
+    dest_zip: Path,
+    mode: Literal["r", "w", "x", "a"] = "w",
+) -> None:
+    if not dest_zip.parent.exists():
+        SecurePath(dest_zip).parent.mkdir(parents=True)
+
+    with ZipFile(dest_zip, mode, ZIP_DEFLATED, allowZip64=True) as package_zip:
+        cli_console.step(f"Creating: {dest_zip}")
+        for src, _ in bundle_map.all_mappings(expand_directories=True):
+            if src.is_file():
+                log.debug("Adding %s to %s", src, dest_zip)
+                package_zip.write(src, arcname=_path_without_top_level_directory(src))
+
+
+def _path_without_top_level_directory(path: Path) -> str:
+    path_parts = path.parts
+    if len(path_parts) > 1:
+        return str(Path(*path_parts[1:]))
+    return str(path)
+
+
 def _to_be_zipped(file: Path) -> bool:
     for pattern in IGNORED_FILES:
         # This has to be a string because of fnmatch