snowflake-cli-labs 2.8.0rc1__py3-none-any.whl → 3.0.0__py3-none-any.whl

snowflake/cli/_plugins/snowpark/snowpark_entity_model.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List, Literal, Optional, Union
+
+from pydantic import Field, field_validator
+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.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):
+    handler: str = Field(
+        title="Function’s or procedure’s implementation of the object inside source module",
+        examples=["functions.hello_function"],
+    )
+    returns: str = Field(
+        title="Type of the result"
+    )  # TODO: again, consider Literal/Enum
+    signature: Union[str, List[Argument]] = Field(
+        title="The signature parameter describes consecutive arguments passed to the object"
+    )
+    runtime: Optional[Union[str, float]] = Field(
+        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)
+            else:
+                _artifacts.append(PathMapping(src=artefact))
+        return _artifacts
+
+    @field_validator("runtime")
+    @classmethod
+    def convert_runtime(cls, runtime_input: Union[str, float]) -> str:
+        if isinstance(runtime_input, float):
+            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)
+
+
+class ProcedureEntityModel(SnowparkEntityModel):
+    type: Literal["procedure"] = DiscriminatorField()  # noqa: A003
+    execute_as_caller: Optional[bool] = Field(
+        title="Determine whether the procedure is executed with the privileges of "
+        "the owner (you) or with the privileges of the caller",
+        default=False,
+    )
+
+
+class FunctionEntityModel(SnowparkEntityModel):
+    type: Literal["function"] = DiscriminatorField()  # noqa: A003
+
+
+class UdfSprocIdentifier:
+    def __init__(self, identifier: FQN, arg_names, arg_types, arg_defaults):
+        self._identifier = identifier
+        self._arg_names = arg_names
+        self._arg_types = arg_types
+        self._arg_defaults = arg_defaults
+
+    def _identifier_from_signature(self, sig: List[str], for_sql: bool = False):
+        signature = self._comma_join(sig)
+        id_ = self._identifier.sql_identifier if for_sql else self._identifier
+        return f"{id_}({signature})"
+
+    @staticmethod
+    def _comma_join(*args):
+        return ", ".join(*args)
+
+    @property
+    def identifier_with_arg_names(self):
+        return self._identifier_from_signature(self._arg_names)
+
+    @property
+    def identifier_with_arg_types(self):
+        return self._identifier_from_signature(self._arg_types)
+
+    @property
+    def identifier_with_arg_names_types(self):
+        sig = [f"{n} {t}" for n, t in zip(self._arg_names, self._arg_types)]
+        return self._identifier_from_signature(sig)
+
+    @property
+    def identifier_with_arg_names_types_defaults(self):
+        return self._identifier_from_signature(self._full_signature())
+
+    def _is_signature_type_a_string(self, sig_type: str) -> bool:
+        return sig_type.lower() in ["string", "varchar"]
+
+    def _full_signature(self):
+        sig = []
+        for name, _type, _default in zip(
+            self._arg_names, self._arg_types, self._arg_defaults
+        ):
+            s = f"{name} {_type}"
+            if _default:
+                if self._is_signature_type_a_string(_type):
+                    _default = f"'{_default}'"
+                s += f" default {_default}"
+            sig.append(s)
+        return sig
+
+    @property
+    def identifier_for_sql(self):
+        return self._identifier_from_signature(self._full_signature(), for_sql=True)
+
+    @classmethod
+    def from_definition(cls, udf_sproc: SnowparkEntityModel):
+        names = []
+        types = []
+        defaults = []
+        if udf_sproc.signature and udf_sproc.signature != "null":
+            for arg in udf_sproc.signature:
+                names.append(arg.name)  # type:ignore
+                types.append(arg.arg_type)  # type:ignore
+                defaults.append(arg.default)  # type:ignore
+
+        identifier = udf_sproc.fqn.using_context()
+        return cls(identifier, names, types, defaults)

snowflake/cli/_plugins/snowpark/snowpark_project_paths.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+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.secure_path import SecurePath
+
+
+@dataclass
+class SnowparkProjectPaths:
+    """
+    This class represents 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)
+
+    @property
+    def snowflake_requirements(self) -> SecurePath:
+        return SecurePath(
+            self.path_relative_to_root(Path("requirements.snowflake.txt"))
+        )
+
+    @property
+    def requirements(self) -> SecurePath:
+        return SecurePath(self.path_relative_to_root(Path("requirements.txt")))
+
+    @property
+    def dependencies(self) -> Path:
+        return self.path_relative_to_root(Path("dependencies.zip"))
+
+
+@dataclass(unsafe_hash=True)
+class Artefact:
+    """Helper for getting paths related to given artefact."""
+
+    path: Path
+    dest: str | None = None
+
+    @property
+    def _artefact_name(self) -> str:
+        if self.path.is_dir():
+            return self.path.stem + ".zip"
+        return self.path.name
+
+    @property
+    def post_build_path(self) -> Path:
+        """
+        Returns post-build artefact path. Directories are mapped to corresponding .zip files.
+        """
+        return self.path.parent / self._artefact_name
+
+    def upload_path(self, stage: FQN | str | None) -> str:
+        """
+        Path on stage to which the artefact 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 = stage_path / self.dest
+        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._artefact_name
+
+    def build(self) -> None:
+        """Build the artefact. Applies only to directories. Files are untouched."""
+        if not self.path.is_dir():
+            return
+        cli_console.step(f"Creating: {self.post_build_path.name}")
+        zip_dir(
+            source=self.path,
+            dest_zip=self.post_build_path,
+        )

snowflake/cli/{plugins → _plugins}/snowpark/snowpark_shared.py

@@ -14,33 +14,7 @@
 
 from __future__ import annotations
 
-import click
 import typer
-from snowflake.cli.api.commands.flags import (
-    deprecated_flag_callback,
-    deprecated_flag_callback_enum,
-)
-from snowflake.cli.plugins.snowpark.models import YesNoAsk
-
-
-def deprecated_allow_native_libraries_option(old_flag_name: str):
-    return typer.Option(
-        YesNoAsk.NO.value,
-        old_flag_name,
-        help="Allows native libraries, when using packages installed through PIP",
-        hidden=True,
-        callback=deprecated_flag_callback_enum(
-            f"{old_flag_name} flag is deprecated. Use --allow-shared-libraries flag instead."
-        ),
-    )
-
-
-def resolve_allow_shared_libraries_yes_no_ask(allow_shared_libraries: YesNoAsk) -> bool:
-    if allow_shared_libraries == YesNoAsk.ASK:
-        return click.confirm("Continue with package installation?", default=False)
-    else:
-        return allow_shared_libraries == YesNoAsk.YES
-
 
 AllowSharedLibrariesOption: bool = typer.Option(
     False,
@@ -48,16 +22,6 @@ AllowSharedLibrariesOption: bool = typer.Option(
     help="Allows shared (.so) libraries, when using packages installed through PIP.",
 )
 
-DeprecatedCheckAnacondaForPyPiDependencies: bool = typer.Option(
-    True,
-    "--check-anaconda-for-pypi-deps/--no-check-anaconda-for-pypi-deps",
-    "-a",
-    help="""Checks if any of missing Anaconda packages dependencies can be imported directly from Anaconda. Valid values include: `true`, `false`, Default: `true`.""",
-    hidden=True,
-    callback=deprecated_flag_callback(
-        "--check-anaconda-for-pypi-deps flag is deprecated. Use --ignore-anaconda flag instead."
-    ),
-)
 
 IgnoreAnacondaOption: bool = typer.Option(
     False,

snowflake/cli/{plugins → _plugins}/snowpark/zipper.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 import fnmatch
 import logging
 from pathlib import Path
-from typing import Iterator, Literal
+from typing import Dict, List, Literal
 from zipfile import ZIP_DEFLATED, ZipFile
 
 log = logging.getLogger(__name__)
@@ -59,16 +59,24 @@ def add_file_to_existing_zip(zip_file: str, file: str):
 
 
 def zip_dir(
-    source: Path, dest_zip: Path, mode: Literal["r", "w", "x", "a"] = "w"
+    source: Path | List[Path],
+    dest_zip: Path,
+    mode: Literal["r", "w", "x", "a"] = "w",
 ) -> None:
-    files_to_pack: Iterator[Path] = filter(
-        _to_be_zipped, map(lambda f: f.absolute(), source.glob("**/*"))
-    )
+
+    if isinstance(source, Path):
+        source = [source]
+
+    files_to_pack: Dict[Path, List[Path]] = {
+        src: list(filter(_to_be_zipped, (f.absolute() for f in src.glob("**/*"))))
+        for src in source
+    }
 
     with ZipFile(dest_zip, mode, ZIP_DEFLATED, allowZip64=True) as package_zip:
-        for file in files_to_pack:
-            log.debug("Adding %s to %s", file, dest_zip)
-            package_zip.write(file, arcname=file.relative_to(source.absolute()))
+        for src, files in files_to_pack.items():
+            for file in files:
+                log.debug("Adding %s to %s", file, dest_zip)
+                package_zip.write(file, arcname=file.relative_to(src))
 
 
 def _to_be_zipped(file: Path) -> bool:

snowflake/cli/{plugins → _plugins}/spcs/__init__.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.plugins.spcs.compute_pool.commands import (
+from snowflake.cli._plugins.spcs.compute_pool.commands import (
     app as compute_pools_app,
 )
-from snowflake.cli.plugins.spcs.image_registry.commands import app as registry_app
-from snowflake.cli.plugins.spcs.image_repository.commands import (
+from snowflake.cli._plugins.spcs.image_registry.commands import app as registry_app
+from snowflake.cli._plugins.spcs.image_repository.commands import (
     app as image_repository_app,
 )
-from snowflake.cli.plugins.spcs.jobs.commands import app as jobs_app
-from snowflake.cli.plugins.spcs.services.commands import app as services_app
+from snowflake.cli._plugins.spcs.services.commands import app as services_app
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 
 app = SnowTyperFactory(
     name="spcs",
@@ -30,6 +29,5 @@ app = SnowTyperFactory(
 
 app.add_typer(compute_pools_app)
 app.add_typer(services_app)
-app.add_typer(jobs_app)
 app.add_typer(registry_app)
 app.add_typer(image_repository_app)

snowflake/cli/{plugins → _plugins}/spcs/compute_pool/commands.py

@@ -18,23 +18,25 @@ from typing import Optional
 
 import typer
 from click import ClickException
+from snowflake.cli._plugins.object.command_aliases import (
+    add_object_command_aliases,
+)
+from snowflake.cli._plugins.object.common import CommentOption
+from snowflake.cli._plugins.spcs.common import (
+    validate_and_set_instances,
+)
+from snowflake.cli._plugins.spcs.compute_pool.manager import ComputePoolManager
 from snowflake.cli.api.commands.flags import (
     IfNotExistsOption,
     OverrideableOption,
+    identifier_argument,
     like_option,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import CommandResult, SingleQueryResult
 from snowflake.cli.api.project.util import is_valid_object_name
-from snowflake.cli.plugins.object.command_aliases import (
-    add_object_command_aliases,
-)
-from snowflake.cli.plugins.object.common import CommentOption
-from snowflake.cli.plugins.spcs.common import (
-    validate_and_set_instances,
-)
-from snowflake.cli.plugins.spcs.compute_pool.manager import ComputePoolManager
 
 app = SnowTyperFactory(
     name="compute-pool",
@@ -43,22 +45,21 @@ app = SnowTyperFactory(
 )
 
 
-def _compute_pool_name_callback(name: str) -> str:
+def _compute_pool_name_callback(name: FQN) -> FQN:
     """
     Verifies that compute pool name is a single valid identifier.
     """
-    if not is_valid_object_name(name, max_depth=0, allow_quoted=False):
+    if not is_valid_object_name(name.identifier, max_depth=0, allow_quoted=False):
         raise ClickException(
             f"'{name}' is not a valid compute pool name. Note that compute pool names must be unquoted identifiers."
         )
     return name
 
 
-ComputePoolNameArgument = typer.Argument(
-    ...,
-    help="Name of the compute pool.",
+ComputePoolNameArgument = identifier_argument(
+    sf_object="compute pool",
+    example="my_compute_pool",
     callback=_compute_pool_name_callback,
-    show_default=False,
 )
 
 
@@ -106,7 +107,7 @@ add_object_command_aliases(
 
 @app.command(requires_connection=True)
 def create(
-    name: str = ComputePoolNameArgument,
+    name: FQN = ComputePoolNameArgument,
     instance_family: str = typer.Option(
         ...,
         "--family",
@@ -131,7 +132,7 @@ def create(
     """
     max_nodes = validate_and_set_instances(min_nodes, max_nodes, "nodes")
     cursor = ComputePoolManager().create(
-        pool_name=name,
+        pool_name=name.identifier,
         min_nodes=min_nodes,
         max_nodes=max_nodes,
         instance_family=instance_family,
@@ -145,33 +146,33 @@ def create(
 
 
 @app.command("stop-all", requires_connection=True)
-def stop_all(name: str = ComputePoolNameArgument, **options) -> CommandResult:
+def stop_all(name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
     """
     Deletes all services running on the compute pool.
     """
-    cursor = ComputePoolManager().stop(pool_name=name)
+    cursor = ComputePoolManager().stop(pool_name=name.identifier)
     return SingleQueryResult(cursor)
 
 
 @app.command(requires_connection=True)
-def suspend(name: str = ComputePoolNameArgument, **options) -> CommandResult:
+def suspend(name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
     """
     Suspends the compute pool by suspending all currently running services and then releasing compute pool nodes.
     """
-    return SingleQueryResult(ComputePoolManager().suspend(name))
+    return SingleQueryResult(ComputePoolManager().suspend(name.identifier))
 
 
 @app.command(requires_connection=True)
-def resume(name: str = ComputePoolNameArgument, **options) -> CommandResult:
+def resume(name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
     """
     Resumes the compute pool from a SUSPENDED state.
     """
-    return SingleQueryResult(ComputePoolManager().resume(name))
+    return SingleQueryResult(ComputePoolManager().resume(name.identifier))
 
 
 @app.command("set", requires_connection=True)
 def set_property(
-    name: str = ComputePoolNameArgument,
+    name: FQN = ComputePoolNameArgument,
     min_nodes: Optional[int] = MinNodesOption(default=None, show_default=False),
     max_nodes: Optional[int] = MaxNodesOption(show_default=False),
     auto_resume: Optional[bool] = AutoResumeOption(default=None, show_default=False),
@@ -187,7 +188,7 @@ def set_property(
     Sets one or more properties for the compute pool.
     """
     cursor = ComputePoolManager().set_property(
-        pool_name=name,
+        pool_name=name.identifier,
         min_nodes=min_nodes,
         max_nodes=max_nodes,
         auto_resume=auto_resume,
@@ -199,7 +200,7 @@ def set_property(
 
 @app.command("unset", requires_connection=True)
 def unset_property(
-    name: str = ComputePoolNameArgument,
+    name: FQN = ComputePoolNameArgument,
     auto_resume: bool = AutoResumeOption(
         default=False,
         param_decls=["--auto-resume"],
@@ -223,7 +224,7 @@ def unset_property(
     Resets one or more properties for the compute pool to their default value(s).
     """
     cursor = ComputePoolManager().unset_property(
-        pool_name=name,
+        pool_name=name.identifier,
         auto_resume=auto_resume,
         auto_suspend_secs=auto_suspend_secs,
         comment=comment,
@@ -232,9 +233,9 @@ def unset_property(
 
 
 @app.command(requires_connection=True)
-def status(pool_name: str = ComputePoolNameArgument, **options) -> CommandResult:
+def status(pool_name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
     """
     Retrieves the status of a compute pool along with a relevant message, if one exists.
     """
-    cursor = ComputePoolManager().status(pool_name=pool_name)
+    cursor = ComputePoolManager().status(pool_name=pool_name.identifier)
     return SingleQueryResult(cursor)

snowflake/cli/{plugins → _plugins}/spcs/compute_pool/manager.py

@@ -16,13 +16,13 @@ from __future__ import annotations
 
 from typing import List, Optional
 
-from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.api.sql_execution import SqlExecutionMixin
-from snowflake.cli.plugins.spcs.common import (
+from snowflake.cli._plugins.spcs.common import (
     NoPropertiesProvidedError,
     handle_object_already_exists,
     strip_empty_lines,
 )
+from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
 

snowflake/cli/{plugins → _plugins}/spcs/image_registry/commands.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.api.output.types import MessageResult, ObjectResult
-from snowflake.cli.plugins.spcs.image_registry.manager import (
+from snowflake.cli._plugins.spcs.image_registry.manager import (
     RegistryManager,
 )
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.output.types import MessageResult, ObjectResult
 
 app = SnowTyperFactory(
     name="image-registry",