snowflake-cli-labs 2.3.1__py3-none-any.whl → 2.4.0__py3-none-any.whl

snowflake/cli/plugins/nativeapp/codegen/snowpark/python_processor.py

@@ -0,0 +1,489 @@
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+from textwrap import dedent
+from typing import Any, Dict, List, Optional, Set
+
+from click import ClickException
+from snowflake.cli.api.console import cli_console as cc
+from snowflake.cli.api.project.errors import SchemaValidationError
+from snowflake.cli.api.project.schemas.native_app.native_app import NativeApp
+from snowflake.cli.api.project.schemas.native_app.path_mapping import (
+    PathMapping,
+    ProcessorMapping,
+)
+from snowflake.cli.api.utils.rendering import jinja_render_from_file
+from snowflake.cli.plugins.nativeapp.artifacts import (
+    BundleMap,
+    find_setup_script_file,
+)
+from snowflake.cli.plugins.nativeapp.codegen.artifact_processor import ArtifactProcessor
+from snowflake.cli.plugins.nativeapp.codegen.sandbox import (
+    ExecutionEnvironmentType,
+    SandboxExecutionError,
+    execute_script_in_sandbox,
+)
+from snowflake.cli.plugins.nativeapp.codegen.snowpark.extension_function_utils import (
+    deannotate_module_source,
+    ensure_all_string_literals,
+    ensure_string_literal,
+    get_function_type_signature_for_grant,
+    get_qualified_object_name,
+    get_sql_argument_signature,
+    get_sql_object_type,
+)
+from snowflake.cli.plugins.nativeapp.codegen.snowpark.models import (
+    ExtensionFunctionTypeEnum,
+    NativeAppExtensionFunction,
+)
+from snowflake.cli.plugins.stage.diff import to_stage_path
+
+DEFAULT_TIMEOUT = 30
+TEMPLATE_PATH = Path(__file__).parent / "callback_source.py.jinja"
+SNOWPARK_LIB_NAME = "snowflake-snowpark-python"
+SNOWPARK_LIB_REGEX = rf"'{SNOWPARK_LIB_NAME}\s*((<|<=|!=|==|>=|>|~=|===)\s*[a-zA-Z0-9_.*+!-]+)?'"  # support PEP 508, even though not all of it is supported in Snowflake yet
+STAGE_PREFIX = "@"
+
+
+def _determine_virtual_env(
+    project_root: Path, processor: ProcessorMapping
+) -> Dict[str, Any]:
+    """
+    Determines a virtual environment to run the Snowpark processor in, either through the project definition or by querying the current environment.
+    """
+    if (processor.properties is None) or ("env" not in processor.properties):
+        return {}
+
+    env_props = processor.properties["env"]
+    env_type = env_props.get("type", None)
+
+    if env_type is None:
+        return {}
+
+    if env_type.upper() == ExecutionEnvironmentType.CONDA.name:
+        env_name = env_props.get("name", None)
+        if env_name is None:
+            cc.warning(
+                "No name found in project definition file for the conda environment to run the Snowpark processor in. Will attempt to auto-detect the current conda environment."
+            )
+        return {"env_type": ExecutionEnvironmentType.CONDA, "name": env_name}
+    elif env_type.upper() == ExecutionEnvironmentType.VENV.name:
+        env_path_str = env_props.get("path", None)
+        if env_path_str is None:
+            cc.warning(
+                "No path found in project definition file for the conda environment to run the Snowpark processor in. Will attempt to auto-detect the current venv path."
+            )
+            env_path = None
+        else:
+            env_path = Path(env_path_str)
+            if not env_path.is_absolute():
+                env_path = project_root / env_path
+        return {
+            "env_type": ExecutionEnvironmentType.VENV,
+            "path": env_path,
+        }
+    elif env_type.upper() == ExecutionEnvironmentType.CURRENT.name:
+        return {
+            "env_type": ExecutionEnvironmentType.CURRENT,
+        }
+    return {}
+
+
+def _is_python_file_artifact(src: Path, dest: Path):
+    return src.is_file() and src.suffix == ".py"
+
+
+def _execute_in_sandbox(
+    py_file: str, deploy_root: Path, kwargs: Dict[str, Any]
+) -> Optional[List[Dict[str, Any]]]:
+    # Create the code snippet to be executed in the sandbox
+    script_source = jinja_render_from_file(
+        template_path=TEMPLATE_PATH, data={"py_file": py_file}
+    )
+
+    try:
+        completed_process = execute_script_in_sandbox(
+            script_source=script_source,
+            cwd=deploy_root,
+            timeout=DEFAULT_TIMEOUT,
+            **kwargs,
+        )
+    except SandboxExecutionError as sdbx_err:
+        cc.warning(
+            f"Could not fetch Snowpark objects from {py_file} due to {sdbx_err}, continuing execution for the rest of the python files."
+        )
+        return None
+    except Exception as err:
+        cc.warning(
+            f"Could not fetch Snowpark objects from {py_file} due to {err}, continuing execution for the rest of the python files."
+        )
+        return None
+
+    if completed_process.returncode != 0:
+        cc.warning(
+            f"Could not fetch Snowpark objects from {py_file} due to the following error:\n {completed_process.stderr}"
+        )
+        cc.warning("Continuing execution for the rest of the python files.")
+        return None
+
+    try:
+        return json.loads(completed_process.stdout)
+    except Exception as exc:
+        cc.warning(
+            f"Could not load JSON into python due to the following exception: {exc}"
+        )
+        cc.warning(f"Continuing execution for the rest of the python files.")
+        return None
+
+
+class SnowparkAnnotationProcessor(ArtifactProcessor):
+    """
+    Built-in Processor to discover Snowpark-annotated objects in a given set of python files,
+    and generate SQL code for creation of extension functions based on those discovered objects.
+    """
+
+    def __init__(
+        self,
+        project_definition: NativeApp,
+        project_root: Path,
+        deploy_root: Path,
+        generated_root: Path,
+    ):
+        super().__init__(
+            project_definition=project_definition,
+            project_root=project_root,
+            deploy_root=deploy_root,
+            generated_root=generated_root,
+        )
+        self.project_definition = project_definition
+        self.project_root = project_root
+        self.deploy_root = deploy_root
+        self.generated_root = generated_root
+
+        if self.generated_root.exists():
+            raise ClickException(
+                f"Path {self.generated_root} already exists. Please choose a different name for your generated directory in the project definition file."
+            )
+
+    def process(
+        self,
+        artifact_to_process: PathMapping,
+        processor_mapping: Optional[ProcessorMapping],
+        **kwargs,
+    ) -> None:
+        """
+        Collects code annotations from Snowpark python files containing extension functions and augments the existing
+        setup script with generated SQL that registers these functions.
+        """
+
+        bundle_map = BundleMap(
+            project_root=self.project_root, deploy_root=self.deploy_root
+        )
+        bundle_map.add(artifact_to_process)
+
+        collected_extension_functions_by_path = self.collect_extension_functions(
+            bundle_map, processor_mapping
+        )
+
+        collected_output = []
+        collected_sql_files: List[Path] = []
+        for py_file, extension_fns in collected_extension_functions_by_path.items():
+            sql_file = self.generate_new_sql_file_name(
+                py_file=py_file,
+            )
+            collected_sql_files.append(sql_file)
+            for extension_fn in extension_fns:
+                create_stmt = generate_create_sql_ddl_statement(extension_fn)
+                if create_stmt is None:
+                    continue
+
+                relative_py_file = py_file.relative_to(bundle_map.deploy_root())
+
+                grant_statements = generate_grant_sql_ddl_statements(extension_fn)
+                if grant_statements is not None:
+                    collected_output.append(grant_statements)
+
+                with open(sql_file, "a") as file:
+                    file.write(
+                        f"-- Generated by the Snowflake CLI from {relative_py_file}\n"
+                    )
+                    file.write(f"-- DO NOT EDIT\n")
+                    file.write(create_stmt)
+                    if grant_statements is not None:
+                        file.write("\n")
+                        file.write(grant_statements)
+
+            self.deannotate(py_file, extension_fns)
+
+        if collected_sql_files:
+            edit_setup_script_with_exec_imm_sql(
+                collected_sql_files=collected_sql_files,
+                deploy_root=bundle_map.deploy_root(),
+                generated_root=self.generated_root,
+            )
+
+    def _normalize_imports(
+        self,
+        extension_fn: NativeAppExtensionFunction,
+        py_file: Path,
+        deploy_root: Path,
+    ):
+        normalized_imports: Set[str] = set()
+        # Add the py_file, which is the source of the extension function
+        normalized_imports.add(f"/{to_stage_path(py_file)}")
+
+        for raw_import in extension_fn.imports:
+            if not Path(deploy_root, raw_import).exists():
+                # This should capture import_str of different forms: stagenames, malformed paths etc
+                # But this will also return True if import_str == "/". Regardless, we append it all to normalized_imports
+                cc.warning(
+                    f"{raw_import} does not exist in the deploy root. Skipping validation of this import."
+                )
+
+            if raw_import.startswith(STAGE_PREFIX) or raw_import.startswith("/"):
+                normalized_imports.add(raw_import)
+            else:
+                normalized_imports.add(f"/{to_stage_path(Path(raw_import))}")
+
+        # To ensure order when running tests
+        sorted_imports = list(normalized_imports)
+        sorted_imports.sort()
+        extension_fn.imports = sorted_imports
+
+    def _normalize(
+        self,
+        extension_fn: NativeAppExtensionFunction,
+        py_file: Path,
+        deploy_root: Path,
+    ):
+        if extension_fn.name is None:
+            # The extension function was not named explicitly, use the name of the Python function object as its name
+            extension_fn.name = extension_fn.handler
+
+        # Compute the fully qualified handler
+        # If user defined their udf as @udf(lambda: x, ...) then extension_fn.handler is <lambda>.
+        extension_fn.handler = f"{py_file.stem}.{extension_fn.handler}"
+
+        snowpark_lib_found = False
+        snowpark_lib_regex = re.compile(SNOWPARK_LIB_REGEX)
+        for pkg in extension_fn.packages:
+            if snowpark_lib_regex.fullmatch(ensure_string_literal(pkg.strip())):
+                snowpark_lib_found = True
+                break
+        if not snowpark_lib_found:
+            extension_fn.packages.append(SNOWPARK_LIB_NAME)
+
+        if extension_fn.imports is None:
+            extension_fn.imports = []
+        self._normalize_imports(
+            extension_fn=extension_fn,
+            py_file=py_file,
+            deploy_root=deploy_root,
+        )
+
+    def collect_extension_functions(
+        self, bundle_map: BundleMap, processor_mapping: Optional[ProcessorMapping]
+    ) -> Dict[Path, List[NativeAppExtensionFunction]]:
+        kwargs = (
+            _determine_virtual_env(self.project_root, processor_mapping)
+            if processor_mapping is not None
+            else {}
+        )
+
+        collected_extension_fns_by_path: Dict[
+            Path, List[NativeAppExtensionFunction]
+        ] = {}
+
+        for src_file, dest_file in bundle_map.all_mappings(
+            absolute=True, expand_directories=True, predicate=_is_python_file_artifact
+        ):
+            cc.step(
+                "Processing Snowpark annotations from {}".format(
+                    dest_file.relative_to(bundle_map.deploy_root())
+                )
+            )
+            collected_extension_function_json = _execute_in_sandbox(
+                py_file=str(dest_file.resolve()),
+                deploy_root=self.deploy_root,
+                kwargs=kwargs,
+            )
+
+            if collected_extension_function_json is None:
+                cc.warning(f"Error processing extension functions in {src_file}")
+                cc.warning("Skipping generating code of all objects from this file.")
+                continue
+
+            collected_extension_functions = []
+            for extension_function_json in collected_extension_function_json:
+                try:
+                    extension_fn = NativeAppExtensionFunction(**extension_function_json)
+                    self._normalize(
+                        extension_fn,
+                        py_file=dest_file.relative_to(bundle_map.deploy_root()),
+                        deploy_root=bundle_map.deploy_root(),
+                    )
+                    collected_extension_functions.append(extension_fn)
+                except SchemaValidationError:
+                    cc.warning("Invalid extension function definition")
+
+            if collected_extension_functions:
+                collected_extension_fns_by_path[
+                    dest_file
+                ] = collected_extension_functions
+
+        return collected_extension_fns_by_path
+
+    def generate_new_sql_file_name(self, py_file: Path) -> Path:
+        """
+        Generates a SQL filename for the generated root from the python file, and creates its parent directories.
+        """
+        relative_py_file = py_file.relative_to(self.deploy_root)
+        sql_file = Path(self.generated_root, relative_py_file.with_suffix(".sql"))
+        if sql_file.exists():
+            cc.warning(
+                f"""\
+                File {sql_file} already exists, will append SQL statements to this file.
+            """
+            )
+        sql_file.parent.mkdir(exist_ok=True, parents=True)
+        return sql_file
+
+    def deannotate(
+        self, py_file: Path, extension_fns: List[NativeAppExtensionFunction]
+    ):
+        with open(py_file, "r", encoding="utf-8") as f:
+            code = f.read()
+
+        if py_file.is_symlink():
+            # if the file is a symlink, make sure we don't overwrite the original
+            py_file.unlink()
+
+        new_code = deannotate_module_source(code, extension_fns)
+
+        with open(py_file, "w", encoding="utf-8") as f:
+            f.write(new_code)
+
+
+def generate_create_sql_ddl_statement(
+    extension_fn: NativeAppExtensionFunction,
+) -> Optional[str]:
+    """
+    Generates a "CREATE FUNCTION/PROCEDURE ... " SQL DDL statement based on an extension function definition.
+    Logic for this create statement has been lifted from snowflake-snowpark-python v1.15.0 package.
+    """
+
+    object_type = get_sql_object_type(extension_fn)
+    if object_type is None:
+        cc.warning(f"Unsupported extension function type: {extension_fn.function_type}")
+        return None
+
+    arguments_in_sql = ", ".join(
+        [get_sql_argument_signature(arg) for arg in extension_fn.signature]
+    )
+
+    create_query = dedent(
+        f"""
+               CREATE OR REPLACE
+               {object_type} {get_qualified_object_name(extension_fn)}({arguments_in_sql})
+               RETURNS {extension_fn.returns}
+               LANGUAGE PYTHON
+               RUNTIME_VERSION={extension_fn.runtime}
+    """
+    ).strip()
+
+    if extension_fn.imports:
+        create_query += (
+            f"\nIMPORTS=({', '.join(ensure_all_string_literals(extension_fn.imports))})"
+        )
+
+    if extension_fn.packages:
+        create_query += f"\nPACKAGES=({', '.join(ensure_all_string_literals([pkg.strip() for pkg in extension_fn.packages]))})"
+
+    if extension_fn.external_access_integrations:
+        create_query += f"\nEXTERNAL_ACCESS_INTEGRATIONS=({', '.join(ensure_all_string_literals(extension_fn.external_access_integrations))})"
+
+    if extension_fn.secrets:
+        create_query += f"""\nSECRETS=({', '.join([f"{ensure_string_literal(k)}={v}" for k, v in extension_fn.secrets.items()])})"""
+
+    create_query += f"\nHANDLER={ensure_string_literal(extension_fn.handler)}"
+
+    if extension_fn.function_type == ExtensionFunctionTypeEnum.PROCEDURE:
+        if extension_fn.execute_as_caller:
+            create_query += f"\nEXECUTE AS CALLER"
+        else:
+            create_query += f"\nEXECUTE AS OWNER"
+    create_query += ";\n"
+
+    return create_query
+
+
+def generate_grant_sql_ddl_statements(
+    extension_fn: NativeAppExtensionFunction,
+) -> Optional[str]:
+    """
+    Generates a "GRANT USAGE TO ... " SQL DDL statement based on a dictionary of extension function properties.
+    If no application roles are present, then the function returns None.
+    """
+
+    if not extension_fn.application_roles:
+        cc.warning(
+            "Skipping generation of 'GRANT USAGE ON ...' SQL statement for this object due to lack of application roles."
+        )
+        return None
+
+    grant_sql_statements = []
+    for app_role in extension_fn.application_roles:
+        grant_sql_statement = dedent(
+            f"""\
+            GRANT USAGE ON {get_sql_object_type(extension_fn)} {get_qualified_object_name(extension_fn)}({get_function_type_signature_for_grant(extension_fn)})
+            TO APPLICATION ROLE {app_role};
+            """
+        ).strip()
+        grant_sql_statements.append(grant_sql_statement)
+
+    return "\n".join(grant_sql_statements)
+
+
+def edit_setup_script_with_exec_imm_sql(
+    collected_sql_files: List[Path], deploy_root: Path, generated_root: Path
+):
+    """
+    Adds an 'execute immediate' to setup script for every SQL file in the map
+    """
+    # Create a __generated.sql in the __generated folder
+    generated_file_path = Path(generated_root, f"{generated_root.stem}.sql")
+    generated_file_path.parent.mkdir(exist_ok=True, parents=True)
+
+    if generated_file_path.exists():
+        cc.warning(
+            f"""\
+            File {generated_file_path} already exists.
+            Could not complete code generation of Snowpark Extension Functions.
+            """
+        )
+        return
+
+    # For every SQL file, add SQL statement 'execute immediate' to __generated.sql script.
+    with open(generated_file_path, "a") as file:
+        for sql_file in collected_sql_files:
+            sql_file_relative_path = sql_file.relative_to(
+                deploy_root
+            )  # Path on stage, without the leading slash
+            file.write(f"EXECUTE IMMEDIATE FROM '/{sql_file_relative_path}';\n")
+
+    # Find the setup script in the deploy root.
+    setup_file_path = find_setup_script_file(deploy_root=deploy_root)
+    with open(setup_file_path, "r", encoding="utf-8") as file:
+        code = file.read()
+    # Unlink to prevent over-writing source file
+    if setup_file_path.is_symlink():
+        setup_file_path.unlink()
+
+    # Write original contents and the execute immediate sql to the setup script
+    generated_file_relative_path = generated_file_path.relative_to(deploy_root)
+    with open(setup_file_path, "w", encoding="utf-8") as file:
+        file.write(code)
+        file.write(f"\nEXECUTE IMMEDIATE FROM '/{generated_file_relative_path}';")
+        file.write(f"\n")

snowflake/cli/plugins/nativeapp/commands.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from pathlib import Path
 from typing import List, Optional
@@ -31,7 +33,6 @@ from snowflake.cli.plugins.nativeapp.teardown_processor import (
 )
 from snowflake.cli.plugins.nativeapp.utils import (
     get_first_paragraph_from_markdown_file,
-    is_tty_interactive,
     shallow_git_clone,
 )
 from snowflake.cli.plugins.nativeapp.version.commands import app as versions_app
@@ -152,7 +153,7 @@ def app_run(
         The command fails if no release directive exists for your Snowflake account for a given application package, which is determined from the project definition file. Default: unset.""",
         is_flag=True,
     ),
-    interactive: Optional[bool] = InteractiveOption,
+    interactive: bool = InteractiveOption,
     force: Optional[bool] = ForceOption,
     **options,
 ) -> CommandResult:
@@ -164,7 +165,7 @@ def app_run(
     is_interactive = False
     if force:
         policy = AllowAlwaysPolicy()
-    elif interactive or is_tty_interactive():
+    elif interactive:
         is_interactive = True
         policy = AskAlwaysPolicy()
     else:
@@ -214,6 +215,12 @@ def app_open(
 @with_project_definition("native_app")
 def app_teardown(
     force: Optional[bool] = ForceOption,
+    cascade: Optional[bool] = typer.Option(
+        None,
+        help=f"""Whether to drop all application objects owned by the application within the account. Default: false.""",
+        show_default=False,
+    ),
+    interactive: bool = InteractiveOption,
     **options,
 ) -> CommandResult:
     """
@@ -223,7 +230,7 @@ def app_teardown(
         project_definition=cli_context.project_definition,
         project_root=cli_context.project_root,
     )
-    processor.process(force)
+    processor.process(interactive, force, cascade)
     return MessageResult(f"Teardown is now complete.")
 
 

snowflake/cli/plugins/nativeapp/common_flags.py

@@ -1,11 +1,18 @@
 import typer
+from snowflake.cli.plugins.nativeapp.utils import is_tty_interactive
+
+
+def interactive_callback(val):
+    if val is None:
+        return is_tty_interactive()
+    return val
+
 
 InteractiveOption = typer.Option(
-    False,
-    "--interactive",
-    "-i",
-    help=f"""When enabled, this option displays prompts even if the standard input and output are not terminal devices. Defaults to unset.""",
-    is_flag=True,
+    None,
+    help=f"""When enabled, this option displays prompts even if the standard input and output are not terminal devices. Defaults to True in an interactive shell environment, and False otherwise.""",
+    callback=interactive_callback,
+    show_default=False,
 )
 
 ForceOption = typer.Option(