snowflake-cli-labs 2.6.1__py3-none-any.whl → 2.7.0__py3-none-any.whl

snowflake/cli/plugins/snowpark/commands.py

@@ -38,9 +38,7 @@ from snowflake.cli.api.constants import (
     DEPLOYMENT_STAGE,
     ObjectType,
 )
-from snowflake.cli.api.exceptions import (
-    SecretsWithoutExternalAccessIntegrationError,
-)
+from snowflake.cli.api.exceptions import SecretsWithoutExternalAccessIntegrationError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import (
     CollectionResult,
@@ -48,6 +46,7 @@ from snowflake.cli.api.output.types import (
     MessageResult,
     SingleQueryResult,
 )
+from snowflake.cli.api.project.project_verification import assert_project_type
 from snowflake.cli.api.project.schemas.snowpark.callable import (
     FunctionSchema,
     ProcedureSchema,
@@ -68,7 +67,8 @@ from snowflake.cli.plugins.object.commands import (
 from snowflake.cli.plugins.object.manager import ObjectManager
 from snowflake.cli.plugins.snowpark import package_utils
 from snowflake.cli.plugins.snowpark.common import (
-    build_udf_sproc_identifier,
+    FunctionOrProcedure,
+    UdfSprocIdentifier,
     check_if_replace_is_required,
 )
 from snowflake.cli.plugins.snowpark.manager import FunctionManager, ProcedureManager
@@ -116,7 +116,7 @@ add_init_command(app, project_type="Snowpark", template="default_snowpark")
 
 
 @app.command("deploy", requires_connection=True)
-@with_project_definition("snowpark")
+@with_project_definition()
 def deploy(
     replace: bool = ReplaceOption(
         help="Replaces procedure or function, even if no detected changes to metadata"
@@ -128,6 +128,9 @@ def deploy(
     By default, if any of the objects exist already the commands will fail unless `--replace` flag is provided.
     All deployed objects use the same artifact which is deployed only once.
     """
+
+    assert_project_type("snowpark")
+
     snowpark = cli_context.project_definition.snowpark
     paths = SnowparkPackagePaths.for_snowpark_project(
         project_root=SecurePath(cli_context.project_root),
@@ -218,7 +221,7 @@ def deploy(
 
 
 def _assert_object_definitions_are_correct(
-    object_type, object_definitions: List[FunctionSchema | ProcedureSchema]
+    object_type, object_definitions: List[FunctionOrProcedure]
 ):
     for definition in object_definitions:
         database = definition.database
@@ -237,14 +240,14 @@ def _assert_object_definitions_are_correct(
 
 def _find_existing_objects(
     object_type: ObjectType,
-    objects: List[Dict],
+    objects: List[FunctionOrProcedure],
     om: ObjectManager,
 ):
     existing_objects = {}
     for object_definition in objects:
-        identifier = build_udf_sproc_identifier(
-            object_definition, om, include_parameter_names=False
-        )
+        identifier = UdfSprocIdentifier.from_definition(
+            object_definition
+        ).identifier_with_arg_types
         try:
             current_state = om.describe(
                 object_type=object_type.value.sf_name,
@@ -293,50 +296,52 @@ def get_app_stage_path(stage_name: Optional[str], project_name: str) -> str:
 def _deploy_single_object(
     manager: FunctionManager | ProcedureManager,
     object_type: ObjectType,
-    object_definition: FunctionSchema | ProcedureSchema,
+    object_definition: FunctionOrProcedure,
     existing_objects: Dict[str, Dict],
     snowflake_dependencies: List[str],
     stage_artifact_path: str,
 ):
-    identifier = build_udf_sproc_identifier(
-        object_definition, manager, include_parameter_names=False
-    )
-    identifier_with_default_values = build_udf_sproc_identifier(
-        object_definition,
-        manager,
-        include_parameter_names=True,
-        include_default_values=True,
+
+    identifiers = UdfSprocIdentifier.from_definition(object_definition)
+
+    log.info(
+        "Deploying %s: %s", object_type, identifiers.identifier_with_arg_names_types
     )
-    log.info("Deploying %s: %s", object_type, identifier_with_default_values)
 
     handler = object_definition.handler
     returns = object_definition.returns
     imports = object_definition.imports
     external_access_integrations = object_definition.external_access_integrations
+    runtime_ver = object_definition.runtime
+    execute_as_caller = None
+    if object_type == ObjectType.PROCEDURE:
+        execute_as_caller = object_definition.execute_as_caller
     replace_object = False
 
-    object_exists = identifier in existing_objects
+    object_exists = identifiers.identifier_with_arg_types in existing_objects
     if object_exists:
         replace_object = check_if_replace_is_required(
             object_type=object_type,
-            current_state=existing_objects[identifier],
+            current_state=existing_objects[identifiers.identifier_with_arg_types],
             handler=handler,
             return_type=returns,
             snowflake_dependencies=snowflake_dependencies,
             external_access_integrations=external_access_integrations,
             imports=imports,
             stage_artifact_file=stage_artifact_path,
+            runtime_ver=runtime_ver,
+            execute_as_caller=execute_as_caller,
         )
 
     if object_exists and not replace_object:
         return {
-            "object": identifier_with_default_values,
+            "object": identifiers.identifier_with_arg_names_types_defaults,
             "type": str(object_type),
             "status": "packages updated",
         }
 
     create_or_replace_kwargs = {
-        "identifier": identifier_with_default_values,
+        "identifier": identifiers,
         "handler": handler,
         "return_type": returns,
         "artifact_file": stage_artifact_path,
@@ -354,7 +359,7 @@ def _deploy_single_object(
 
     status = "created" if not object_exists else "definition updated"
     return {
-        "object": identifier_with_default_values,
+        "object": identifiers.identifier_with_arg_names_types_defaults,
         "type": str(object_type),
         "status": status,
     }
@@ -379,7 +384,7 @@ def _read_snowflake_requrements_file(file_path: SecurePath):
 
 
 @app.command("build", requires_connection=True)
-@with_project_definition("snowpark")
+@with_project_definition()
 def build(
     ignore_anaconda: bool = IgnoreAnacondaOption,
     allow_shared_libraries: bool = AllowSharedLibrariesOption,
@@ -396,6 +401,9 @@ def build(
     Builds the Snowpark project as a `.zip` archive that can be used by `deploy` command.
     The archive is built using only the `src` directory specified in the project file.
     """
+
+    assert_project_type("snowpark")
+
     if not deprecated_check_anaconda_for_pypi_deps:
         ignore_anaconda = True
     snowpark_paths = SnowparkPackagePaths.for_snowpark_project(

snowflake/cli/plugins/snowpark/common.py

@@ -15,11 +15,14 @@
 from __future__ import annotations
 
 import re
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, Union
 
 from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.identifiers import FQN
-from snowflake.cli.api.project.schemas.snowpark.argument import Argument
+from snowflake.cli.api.project.schemas.snowpark.callable import (
+    FunctionSchema,
+    ProcedureSchema,
+)
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.cli.plugins.snowpark.models import Requirement
 from snowflake.cli.plugins.snowpark.package_utils import (
@@ -28,6 +31,7 @@ from snowflake.cli.plugins.snowpark.package_utils import (
 from snowflake.connector.cursor import SnowflakeCursor
 
 DEFAULT_RUNTIME = "3.8"
+FunctionOrProcedure = Union[FunctionSchema, ProcedureSchema]
 
 
 def check_if_replace_is_required(
@@ -39,6 +43,8 @@ def check_if_replace_is_required(
     external_access_integrations: List[str],
     imports: List[str],
     stage_artifact_file: str,
+    runtime_ver: Optional[str] = None,
+    execute_as_caller: Optional[bool] = None,
 ) -> bool:
     import logging
 
@@ -78,6 +84,22 @@ def check_if_replace_is_required(
         return True
 
     if _compare_imports(resource_json, imports, stage_artifact_file):
+        log.info("Imports do not match. Replacing the %s", object_type)
+        return True
+
+    if runtime_ver is not None and runtime_ver != resource_json.get(
+        "runtime_version", "RUNTIME_NOT_SET"
+    ):
+        log.info("Runtime versions do not match. Replacing the %s", object_type)
+        return True
+
+    if execute_as_caller is not None and (
+        resource_json.get("execute as", "OWNER")
+        != ("CALLER" if execute_as_caller else "OWNER")
+    ):
+        log.info(
+            "Execute as caller settings do not match. Replacing the %s", object_type
+        )
         return True
 
     return False
@@ -150,7 +172,7 @@ class SnowparkObjectManager(SqlExecutionMixin):
 
     def create_query(
         self,
-        identifier: str,
+        identifier: UdfSprocIdentifier,
         return_type: str,
         handler: str,
         artifact_file: str,
@@ -166,7 +188,7 @@ class SnowparkObjectManager(SqlExecutionMixin):
         packages_list = ",".join(f"'{p}'" for p in packages)
 
         query = [
-            f"create or replace {self._object_type.value.sf_name} {identifier}",
+            f"create or replace {self._object_type.value.sf_name} {identifier.identifier_for_sql}",
             f"copy grants",
             f"returns {return_type}",
             "language python",
@@ -198,30 +220,69 @@ def _is_signature_type_a_string(sig_type: str) -> bool:
     return sig_type.lower() in ["string", "varchar"]
 
 
-def build_udf_sproc_identifier(
-    udf_sproc,
-    slq_exec_mixin,
-    include_parameter_names,
-    include_default_values=False,
-):
-    def format_arg(arg: Argument):
-        result = f"{arg.arg_type}"
-        if include_parameter_names:
-            result = f"{arg.name} {result}"
-        if include_default_values and arg.default:
-            val = f"{arg.default}"
-            if _is_signature_type_a_string(arg.arg_type):
-                val = f"'{val}'"
-            result += f" default {val}"
-        return result
-
-    if udf_sproc.signature and udf_sproc.signature != "null":
-        arguments = ", ".join(format_arg(arg) for arg in udf_sproc.signature)
-    else:
-        arguments = ""
+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})"
 
-    name = FQN.from_identifier_model(udf_sproc).using_context().identifier
-    return f"{name}({arguments})"
+    @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 _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 _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: FunctionOrProcedure):
+        names = []
+        types = []
+        defaults = []
+        if udf_sproc.signature and udf_sproc.signature != "null":
+            for arg in udf_sproc.signature:
+                names.append(arg.name)
+                types.append(arg.arg_type)
+                defaults.append(arg.default)
+
+        identifier = FQN.from_identifier_model(udf_sproc).using_context()
+        return cls(identifier, names, types, defaults)
 
 
 def _compare_imports(

snowflake/cli/plugins/snowpark/manager.py

@@ -18,7 +18,10 @@ import logging
 from typing import Dict, List, Optional
 
 from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.plugins.snowpark.common import SnowparkObjectManager
+from snowflake.cli.plugins.snowpark.common import (
+    SnowparkObjectManager,
+    UdfSprocIdentifier,
+)
 from snowflake.connector.cursor import SnowflakeCursor
 
 log = logging.getLogger(__name__)
@@ -35,7 +38,7 @@ class FunctionManager(SnowparkObjectManager):
 
     def create_or_replace(
         self,
-        identifier: str,
+        identifier: UdfSprocIdentifier,
         return_type: str,
         handler: str,
         artifact_file: str,
@@ -45,7 +48,11 @@ class FunctionManager(SnowparkObjectManager):
         secrets: Optional[Dict[str, str]] = None,
         runtime: Optional[str] = None,
     ) -> SnowflakeCursor:
-        log.debug("Creating function %s using @%s", identifier, artifact_file)
+        log.debug(
+            "Creating function %s using @%s",
+            identifier.identifier_with_arg_names_types_defaults,
+            artifact_file,
+        )
         query = self.create_query(
             identifier,
             return_type,
@@ -71,7 +78,7 @@ class ProcedureManager(SnowparkObjectManager):
 
     def create_or_replace(
         self,
-        identifier: str,
+        identifier: UdfSprocIdentifier,
         return_type: str,
         handler: str,
         artifact_file: str,
@@ -82,7 +89,11 @@ class ProcedureManager(SnowparkObjectManager):
         runtime: Optional[str] = None,
         execute_as_caller: bool = False,
     ) -> SnowflakeCursor:
-        log.debug("Creating procedure %s using @%s", identifier, artifact_file)
+        log.debug(
+            "Creating procedure %s using @%s",
+            identifier.identifier_with_arg_names_types_defaults,
+            artifact_file,
+        )
         query = self.create_query(
             identifier,
             return_type,

snowflake/cli/plugins/snowpark/models.py

@@ -33,6 +33,10 @@ class YesNoAsk(Enum):
 class Requirement(requirement.Requirement):
     extra_pattern = re.compile("'([^']*)'")
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.package_name = None
+
     @classmethod
     def parse_line(cls, line: str) -> Requirement:
         if len(line_elements := line.split(";")) > 1:
@@ -44,6 +48,8 @@ class Requirement(requirement.Requirement):
                 if "extra" in element and (extras := cls.extra_pattern.search(element)):
                     result.extras.extend(extras.groups())
 
+        result.package_name = result.name
+
         if result.uri and not result.name:
             result.name = get_package_name(result.uri)
         result.name = cls.standardize_name(result.name)

snowflake/cli/plugins/sql/commands.py

@@ -21,6 +21,7 @@ import typer
 from snowflake.cli.api.commands.decorators import with_project_definition
 from snowflake.cli.api.commands.flags import (
     parse_key_value_variables,
+    variables_option,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.output.types import CommandResult, MultipleResults, QueryResult
@@ -55,11 +56,8 @@ def execute_sql(
         "-i",
         help="Read the query from standard input. Use it when piping input to this command.",
     ),
-    data_override: List[str] = typer.Option(
-        None,
-        "--variable",
-        "-D",
-        help="String in format of key=value. If provided the SQL content will "
+    data_override: List[str] = variables_option(
+        "String in format of key=value. If provided the SQL content will "
         "be treated as template and rendered using provided data.",
     ),
     **options,

snowflake/cli/plugins/sql/manager.py

@@ -22,9 +22,9 @@ from typing import Dict, Iterable, List, Tuple
 
 from click import ClickException, UsageError
 from jinja2 import UndefinedError
+from snowflake.cli.api.rendering.sql_templates import snowflake_sql_jinja_render
 from snowflake.cli.api.secure_path import UNLIMITED, SecurePath
 from snowflake.cli.api.sql_execution import SqlExecutionMixin, VerboseCursor
-from snowflake.cli.api.utils.rendering import snowflake_sql_jinja_render
 from snowflake.cli.plugins.sql.snowsql_templating import transpile_snowsql_templates
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.util_text import split_statements

snowflake/cli/plugins/stage/commands.py

@@ -23,9 +23,9 @@ import click
 import typer
 from snowflake.cli.api.cli_global_context import cli_context
 from snowflake.cli.api.commands.flags import (
+    ExecuteVariablesOption,
     OnErrorOption,
     PatternOption,
-    VariablesOption,
     like_option,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
@@ -201,7 +201,7 @@ def execute(
         show_default=False,
     ),
     on_error: OnErrorType = OnErrorOption,
-    variables: Optional[List[str]] = VariablesOption,
+    variables: Optional[List[str]] = ExecuteVariablesOption,
     **options,
 ):
     """

snowflake/cli/plugins/stage/diff.py

@@ -14,9 +14,7 @@
 
 from __future__ import annotations
 
-import hashlib
 import logging
-import re
 from dataclasses import dataclass, field
 from pathlib import Path, PurePosixPath
 from typing import Collection, Dict, List, Optional, Tuple
@@ -25,14 +23,11 @@ from snowflake.cli.api.console import cli_console as cc
 from snowflake.cli.api.exceptions import (
     SnowflakeSQLExecutionError,
 )
-from snowflake.cli.api.secure_path import UNLIMITED, SecurePath
 from snowflake.cli.plugins.nativeapp.artifacts import BundleMap
 from snowflake.connector.cursor import DictCursor
 
 from .manager import StageManager
-
-MD5SUM_REGEX = r"^[A-Fa-f0-9]{32}$"
-CHUNK_SIZE_BYTES = 8192
+from .md5 import UnknownMD5FormatError, file_matches_md5sum
 
 log = logging.getLogger(__name__)
 
@@ -72,43 +67,6 @@ class DiffResult:
         }
 
 
-def is_valid_md5sum(checksum: str) -> bool:
-    """
-    Could the provided hexadecimal checksum represent a valid md5sum?
-    """
-    return re.match(MD5SUM_REGEX, checksum) is not None
-
-
-def compute_md5sum(file: Path) -> str:
-    """
-    Returns a hexadecimal checksum for the file located at the given path.
-    """
-    if not file.is_file():
-        raise ValueError(
-            "The provided file does not exist or not a (symlink to a) regular file"
-        )
-
-    # FIXME: there are two cases in which this will fail to provide a matching
-    # md5sum, even when the underlying file is the same:
-    #  1. when the stage uses SNOWFLAKE_FULL encryption
-    #  2. when the file was uploaded in multiple parts
-
-    # We can re-create the second if we know what chunk size was used by the
-    # upload process to the backing object store (e.g. S3, azure blob, etc.)
-    # but we cannot re-create the first as the encrpytion key is hidden.
-
-    # We are assuming that we will not get accidental collisions here due to the
-    # large space of the md5sum (32 * 4 = 128 bits means 1-in-9-trillion chance)
-    # combined with the fact that the file name + path must also match elsewhere.
-
-    with SecurePath(file).open("rb", read_file_limit_mb=UNLIMITED) as f:
-        file_hash = hashlib.md5()
-        while chunk := f.read(CHUNK_SIZE_BYTES):
-            file_hash.update(chunk)
-
-    return file_hash.hexdigest()
-
-
 def enumerate_files(path: Path) -> List[Path]:
     """
     Get a list of all files in a directory (recursively).
@@ -131,7 +89,7 @@ def strip_stage_name(path: str) -> StagePath:
     return StagePath(*path.split("/")[1:])
 
 
-def build_md5_map(list_stage_cursor: DictCursor) -> Dict[StagePath, str]:
+def build_md5_map(list_stage_cursor: DictCursor) -> Dict[StagePath, Optional[str]]:
     """
     Returns a mapping of relative stage paths to their md5sums.
     """
@@ -173,30 +131,35 @@ def compute_stage_diff(
 
     for local_file in local_files:
         relpath = local_file.relative_to(local_root)
-        stage_filename = to_stage_path(relpath)
-        if stage_filename not in remote_md5:
+        stage_path = to_stage_path(relpath)
+        if stage_path not in remote_md5:
             # doesn't exist on the stage
-            result.only_local.append(stage_filename)
+            result.only_local.append(stage_path)
         else:
-            # N.B. we could compare local size vs remote size to skip the relatively-
-            # expensive md5sum operation, but after seeing a comment that says the value
-            # may not always be correctly populated, we'll ignore that column.
-            stage_md5sum = remote_md5[stage_filename]
-            if is_valid_md5sum(stage_md5sum) and stage_md5sum == compute_md5sum(
-                local_file
-            ):
-                # the file definitely hasn't changed
-                result.identical.append(stage_filename)
-            else:
-                # either the file has changed, or we can't tell if it has
-                result.different.append(stage_filename)
+            # N.B. file size on stage is not always accurate, so cannot fail fast
+            try:
+                if file_matches_md5sum(local_file, remote_md5[stage_path]):
+                    # We are assuming that we will not get accidental collisions here due to the
+                    # large space of the md5sum (32 * 4 = 128 bits means 1-in-9-trillion chance)
+                    # combined with the fact that the file name + path must also match elsewhere.
+                    result.identical.append(stage_path)
+                else:
+                    # either the file has changed, or we can't tell if it has
+                    result.different.append(stage_path)
+            except UnknownMD5FormatError:
+                log.warning(
+                    "Could not compare md5 for %s, assuming file has changed",
+                    local_file,
+                    exc_info=True,
+                )
+                result.different.append(stage_path)
 
             # mark this file as seen
-            del remote_md5[stage_filename]
+            del remote_md5[stage_path]
 
     # every entry here is a file we never saw locally
-    for stage_filename in remote_md5.keys():
-        result.only_on_stage.append(stage_filename)
+    for stage_path in remote_md5.keys():
+        result.only_on_stage.append(stage_path)
 
     return result
 
@@ -230,8 +193,8 @@ def delete_only_on_stage_files(
     """
     Deletes all files from a Snowflake stage according to the input list of filenames, using a custom role.
     """
-    for _stage_filename in only_on_stage:
-        stage_manager.remove(stage_name=stage_fqn, path=str(_stage_filename), role=role)
+    for _stage_path in only_on_stage:
+        stage_manager.remove(stage_name=stage_fqn, path=str(_stage_path), role=role)
 
 
 def put_files_on_stage(