snowflake-cli-labs 2.6.0rc0__py3-none-any.whl → 2.7.0rc0__py3-none-any.whl

snowflake/cli/plugins/nativeapp/codegen/compiler.py

@@ -24,12 +24,22 @@ from snowflake.cli.plugins.nativeapp.codegen.artifact_processor import (
     ArtifactProcessor,
     UnsupportedArtifactProcessorError,
 )
+from snowflake.cli.plugins.nativeapp.codegen.setup.native_app_setup_processor import (
+    NativeAppSetupProcessor,
+)
 from snowflake.cli.plugins.nativeapp.codegen.snowpark.python_processor import (
     SnowparkAnnotationProcessor,
 )
+from snowflake.cli.plugins.nativeapp.feature_flags import FeatureFlag
 from snowflake.cli.plugins.nativeapp.project_model import NativeAppProjectModel
 
 SNOWPARK_PROCESSOR = "snowpark"
+NA_SETUP_PROCESSOR = "native-app-setup"
+
+_REGISTERED_PROCESSORS_BY_NAME = {
+    SNOWPARK_PROCESSOR: SnowparkAnnotationProcessor,
+    NA_SETUP_PROCESSOR: NativeAppSetupProcessor,
+}
 
 
 class NativeAppCompiler:
@@ -54,28 +64,31 @@ class NativeAppCompiler:
         Go through every artifact object in the project definition of a native app, and execute processors in order of specification for each of the artifact object.
         May have side-effects on the filesystem by either directly editing source files or the deploy root.
         """
-        should_proceed = False
-        for artifact in self._na_project.artifacts:
-            if artifact.processors:
-                should_proceed = True
-                break
-        if not should_proceed:
+
+        if not self._should_invoke_processors():
             return
 
         with cc.phase("Invoking artifact processors"):
+            if self._na_project.generated_root.exists():
+                raise ClickException(
+                    f"Path {self._na_project.generated_root} already exists. Please choose a different name for your generated directory in the project definition file."
+                )
+
             for artifact in self._na_project.artifacts:
                 for processor in artifact.processors:
-                    artifact_processor = self._try_create_processor(
-                        processor_mapping=processor,
-                    )
-                    if artifact_processor is None:
-                        raise UnsupportedArtifactProcessorError(
-                            processor_name=processor.name
-                        )
-                    else:
-                        artifact_processor.process(
-                            artifact_to_process=artifact, processor_mapping=processor
+                    if self._is_enabled(processor):
+                        artifact_processor = self._try_create_processor(
+                            processor_mapping=processor,
                         )
+                        if artifact_processor is None:
+                            raise UnsupportedArtifactProcessorError(
+                                processor_name=processor.name
+                            )
+                        else:
+                            artifact_processor.process(
+                                artifact_to_process=artifact,
+                                processor_mapping=processor,
+                            )
 
     def _try_create_processor(
         self,
@@ -86,15 +99,32 @@ class NativeAppCompiler:
         Fetch processor object if one already exists in the cached_processors dictionary.
         Else, initialize a new object to return, and add it to the cached_processors dictionary.
         """
-        if processor_mapping.name.lower() == SNOWPARK_PROCESSOR:
-            curr_processor = self.cached_processors.get(SNOWPARK_PROCESSOR, None)
-            if curr_processor is not None:
-                return curr_processor
-            else:
-                curr_processor = SnowparkAnnotationProcessor(
-                    na_project=self._na_project,
-                )
-                self.cached_processors[SNOWPARK_PROCESSOR] = curr_processor
-                return curr_processor
-        else:
+        processor_name = processor_mapping.name.lower()
+        current_processor = self.cached_processors.get(processor_name)
+
+        if current_processor is not None:
+            return current_processor
+
+        processor_factory = _REGISTERED_PROCESSORS_BY_NAME.get(processor_name)
+        if processor_factory is None:
+            # No registered processor with the specified name
             return None
+
+        current_processor = processor_factory(
+            na_project=self._na_project,
+        )
+        self.cached_processors[processor_name] = current_processor
+
+        return current_processor
+
+    def _should_invoke_processors(self):
+        for artifact in self._na_project.artifacts:
+            for processor in artifact.processors:
+                if self._is_enabled(processor):
+                    return True
+        return False
+
+    def _is_enabled(self, processor: ProcessorMapping) -> bool:
+        if processor.name.lower() == NA_SETUP_PROCESSOR:
+            return FeatureFlag.ENABLE_NATIVE_APP_PYTHON_SETUP.is_enabled()
+        return True

snowflake/cli/plugins/nativeapp/codegen/sandbox.py

@@ -20,10 +20,13 @@ import subprocess
 import sys
 from enum import Enum
 from pathlib import Path
-from typing import Optional, Sequence, Union
+from typing import Any, Mapping, Optional, Sequence, Union
+from venv import EnvBuilder
 
 from click.exceptions import ClickException
 
+EnvVars = Mapping[str, str]  # Only support str -> str for cross-platform compatibility
+
 
 class SandboxExecutionError(ClickException):
     """An error occurred while executing a python script."""
@@ -57,6 +60,7 @@ def _execute_python_interpreter(
     script_source: str,
     cwd: Optional[Union[str, Path]],
     timeout: Optional[int],
+    env_vars: Optional[EnvVars],
 ) -> subprocess.CompletedProcess:
     if not python_executable:
         raise SandboxExecutionError("No python executable found")
@@ -73,6 +77,7 @@ def _execute_python_interpreter(
         input=script_source,
         timeout=timeout,
         cwd=cwd,
+        env=env_vars,
     )
 
 
@@ -81,6 +86,7 @@ def _execute_in_venv(
     venv_path: Optional[Union[str, Path]] = None,
     cwd: Optional[Union[str, Path]] = None,
     timeout: Optional[int] = None,
+    env_vars: Optional[EnvVars] = None,
 ) -> subprocess.CompletedProcess:
     resolved_venv_path = None
     if venv_path is None:
@@ -114,7 +120,7 @@ def _execute_in_venv(
         )
 
     return _execute_python_interpreter(
-        python_executable, script_source, timeout=timeout, cwd=cwd
+        python_executable, script_source, timeout=timeout, cwd=cwd, env_vars=env_vars
     )
 
 
@@ -123,6 +129,7 @@ def _execute_in_conda_env(
     env_name: Optional[str] = None,
     cwd: Optional[Union[str, Path]] = None,
     timeout: Optional[int] = None,
+    env_vars: Optional[EnvVars] = None,
 ) -> subprocess.CompletedProcess:
     conda_env = env_name
     if conda_env is None:
@@ -142,6 +149,7 @@ def _execute_in_conda_env(
         script_source,
         timeout=timeout,
         cwd=cwd,
+        env_vars=env_vars,
     )
 
 
@@ -149,13 +157,18 @@ def _execute_with_system_path_python(
     script_source: str,
     cwd: Optional[Union[str, Path]] = None,
     timeout: Optional[int] = None,
+    env_vars: Optional[EnvVars] = None,
 ) -> subprocess.CompletedProcess:
     python_executable = (
         shutil.which("python3") or shutil.which("python") or sys.executable
     )
 
     return _execute_python_interpreter(
-        python_executable, script_source, timeout=timeout, cwd=cwd
+        python_executable,
+        script_source,
+        timeout=timeout,
+        cwd=cwd,
+        env_vars=env_vars,
     )
 
 
@@ -172,6 +185,7 @@ def execute_script_in_sandbox(
     env_type: ExecutionEnvironmentType = ExecutionEnvironmentType.AUTO_DETECT,
     cwd: Optional[Union[str, Path]] = None,
     timeout: Optional[int] = None,
+    env_vars: Optional[EnvVars] = None,
     **kwargs,
 ) -> subprocess.CompletedProcess:
     """
@@ -194,24 +208,99 @@ def execute_script_in_sandbox(
     """
     if env_type == ExecutionEnvironmentType.AUTO_DETECT:
         if _is_venv_active():
-            return _execute_in_venv(script_source, cwd=cwd, timeout=timeout)
+            return _execute_in_venv(
+                script_source, cwd=cwd, timeout=timeout, env_vars=env_vars
+            )
         elif _is_conda_active():
-            return _execute_in_conda_env(script_source, cwd=cwd, timeout=timeout)
+            return _execute_in_conda_env(
+                script_source, cwd=cwd, timeout=timeout, env_vars=env_vars
+            )
         else:
             return _execute_with_system_path_python(
-                script_source, cwd=cwd, timeout=timeout
+                script_source, cwd=cwd, timeout=timeout, env_vars=env_vars
             )
     elif env_type == ExecutionEnvironmentType.VENV:
         return _execute_in_venv(
-            script_source, kwargs.get("path"), cwd=cwd, timeout=timeout
+            script_source,
+            kwargs.get("path"),
+            cwd=cwd,
+            timeout=timeout,
+            env_vars=env_vars,
         )
     elif env_type == ExecutionEnvironmentType.CONDA:
         return _execute_in_conda_env(
-            script_source, kwargs.get("name"), cwd=cwd, timeout=timeout
+            script_source,
+            kwargs.get("name"),
+            cwd=cwd,
+            timeout=timeout,
+            env_vars=env_vars,
         )
     elif env_type == ExecutionEnvironmentType.SYSTEM_PATH:
-        return _execute_with_system_path_python(script_source, cwd=cwd, timeout=timeout)
+        return _execute_with_system_path_python(
+            script_source, cwd=cwd, timeout=timeout, env_vars=env_vars
+        )
     else:  # ExecutionEnvironmentType.CURRENT
         return _execute_python_interpreter(
-            sys.executable, script_source, cwd=cwd, timeout=timeout
+            sys.executable, script_source, cwd=cwd, timeout=timeout, env_vars=env_vars
         )
+
+
+class SandboxEnvBuilder(EnvBuilder):
+    """
+    A virtual environment builder that can be used to build an environment suitable for
+    executing user-provided python scripts in an isolated sandbox.
+    """
+
+    def __init__(self, path: Path, **kwargs) -> None:
+        """
+        Creates a new builder with the specified destination path. The path need not
+        exist, it will be created when needed (recursively if necessary).
+
+        Parameters:
+            path (Path): The directory in which the sandbox environment will be created.
+        """
+        super().__init__(**kwargs)
+        self.path = path
+        self._context: Any = None  # cached context
+
+    def post_setup(self, context) -> None:
+        self._context = context
+
+    def ensure_created(self) -> None:
+        """
+        Ensures that the sandbox environment has been created and correctly initialized.
+        """
+        if self.path.exists():
+            self._context = self.ensure_directories(self.path)
+        else:
+            self.path.mkdir(parents=True, exist_ok=True)
+            self.create(
+                self.path
+            )  # will set self._context through the post_setup callback
+
+    def run_python(self, *args) -> str:
+        """
+        Executes the python interpreter in the sandboxed environment with the provided arguments.
+        This raises a CalledProcessError if the python interpreter was not executed successfully.
+
+        Returns:
+            The output of running the command.
+        """
+        positional_args = [
+            self._context.env_exe,
+            "-E",  # passing -E ignores all PYTHON* env vars
+            *args,
+        ]
+        kwargs = {
+            "cwd": self._context.env_dir,
+            "stderr": subprocess.STDOUT,
+        }
+        env = dict(os.environ)
+        env["VIRTUAL_ENV"] = self._context.env_dir
+        return subprocess.check_output(positional_args, **kwargs)
+
+    def pip_install(self, *args: Any) -> None:
+        """
+        Invokes pip install with the provided arguments.
+        """
+        self.run_python("-m", "pip", "install", *[str(arg) for arg in args])

snowflake/cli/plugins/nativeapp/codegen/setup/native_app_setup_processor.py

@@ -0,0 +1,172 @@
+# 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
+
+import json
+import os.path
+from pathlib import Path
+from typing import List, Optional
+
+from click import ClickException
+from snowflake.cli.api.console import cli_console as cc
+from snowflake.cli.api.project.schemas.native_app.path_mapping import (
+    PathMapping,
+    ProcessorMapping,
+)
+from snowflake.cli.plugins.nativeapp.artifacts import BundleMap, find_setup_script_file
+from snowflake.cli.plugins.nativeapp.codegen.artifact_processor import (
+    ArtifactProcessor,
+    is_python_file_artifact,
+)
+from snowflake.cli.plugins.nativeapp.codegen.sandbox import (
+    ExecutionEnvironmentType,
+    SandboxEnvBuilder,
+    execute_script_in_sandbox,
+)
+from snowflake.cli.plugins.nativeapp.project_model import NativeAppProjectModel
+from snowflake.cli.plugins.stage.diff import to_stage_path
+
+DEFAULT_TIMEOUT = 30
+DRIVER_PATH = Path(__file__).parent / "setup_driver.py.source"
+
+
+class NativeAppSetupProcessor(ArtifactProcessor):
+    def __init__(
+        self,
+        na_project: NativeAppProjectModel,
+    ):
+        super().__init__(na_project=na_project)
+
+    def process(
+        self,
+        artifact_to_process: PathMapping,
+        processor_mapping: Optional[ProcessorMapping],
+        **kwargs,
+    ) -> None:
+        """
+        Processes a Python setup script and generates the corresponding SQL commands.
+        """
+        bundle_map = BundleMap(
+            project_root=self._na_project.project_root,
+            deploy_root=self._na_project.deploy_root,
+        )
+        bundle_map.add(artifact_to_process)
+
+        self._create_or_update_sandbox()
+
+        cc.phase("Processing Python setup files")
+
+        files_to_process = []
+        for src_file, dest_file in bundle_map.all_mappings(
+            absolute=True, expand_directories=True, predicate=is_python_file_artifact
+        ):
+            cc.message(
+                f"Found Python setup file: {src_file.relative_to(self._na_project.project_root)}"
+            )
+            files_to_process.append(src_file)
+
+        sql_files_mapping = self._execute_in_sandbox(files_to_process)
+        self._generate_setup_sql(sql_files_mapping)
+
+    def _execute_in_sandbox(self, py_files: List[Path]) -> dict:
+        file_count = len(py_files)
+        cc.step(f"Processing {file_count} setup file{'s' if file_count > 1 else ''}")
+
+        env_vars = {
+            "_SNOWFLAKE_CLI_PROJECT_PATH": str(self._na_project.project_root),
+            "_SNOWFLAKE_CLI_SETUP_FILES": os.pathsep.join(map(str, py_files)),
+            "_SNOWFLAKE_CLI_APP_NAME": str(self._na_project.package_name),
+            "_SNOWFLAKE_CLI_SQL_DEST_DIR": str(self.generated_root),
+        }
+
+        try:
+            result = execute_script_in_sandbox(
+                script_source=DRIVER_PATH.read_text(),
+                env_type=ExecutionEnvironmentType.VENV,
+                cwd=self._na_project.bundle_root,
+                timeout=DEFAULT_TIMEOUT,
+                path=self.sandbox_root,
+                env_vars=env_vars,
+            )
+        except Exception as e:
+            raise ClickException(
+                f"Exception while executing python setup script logic: {e}"
+            )
+
+        if result.returncode == 0:
+            sql_file_mappings = json.loads(result.stdout)
+            return sql_file_mappings
+        else:
+            raise ClickException(
+                f"Failed to execute python setup script logic: {result.stderr}"
+            )
+
+    def _generate_setup_sql(self, sql_file_mappings: dict) -> None:
+        if not sql_file_mappings:
+            # Nothing to generate
+            return
+
+        generated_root = self.generated_root
+        generated_root.mkdir(exist_ok=True, parents=True)
+
+        cc.step("Patching setup script")
+        setup_file_path = find_setup_script_file(
+            deploy_root=self._na_project.deploy_root
+        )
+        with self.edit_file(setup_file_path) as f:
+            new_contents = [f.contents]
+
+            if sql_file_mappings["schemas"]:
+                schemas_file = generated_root / sql_file_mappings["schemas"]
+                new_contents.insert(
+                    0,
+                    f"EXECUTE IMMEDIATE FROM '/{to_stage_path(schemas_file.relative_to(self._na_project.deploy_root))}';",
+                )
+
+            if sql_file_mappings["compute_pools"]:
+                compute_pools_file = generated_root / sql_file_mappings["compute_pools"]
+                new_contents.append(
+                    f"EXECUTE IMMEDIATE FROM '/{to_stage_path(compute_pools_file.relative_to(self._na_project.deploy_root))}';"
+                )
+
+            if sql_file_mappings["services"]:
+                services_file = generated_root / sql_file_mappings["services"]
+                new_contents.append(
+                    f"EXECUTE IMMEDIATE FROM '/{to_stage_path(services_file.relative_to(self._na_project.deploy_root))}';"
+                )
+
+            f.edited_contents = "\n".join(new_contents)
+
+    @property
+    def sandbox_root(self):
+        return self._na_project.bundle_root / "setup_py_venv"
+
+    @property
+    def generated_root(self):
+        return self._na_project.generated_root / "setup_py"
+
+    def _create_or_update_sandbox(self):
+        sandbox_root = self.sandbox_root
+        env_builder = SandboxEnvBuilder(sandbox_root, with_pip=True)
+        if sandbox_root.exists():
+            cc.step("Virtual environment found")
+        else:
+            cc.step(
+                f"Creating virtual environment in {sandbox_root.relative_to(self._na_project.project_root)}"
+            )
+        env_builder.ensure_created()
+
+        # Temporarily fetch the library from a location specified via env vars
+        env_builder.pip_install(os.environ["SNOWFLAKE_APP_PYTHON_LOC"])

snowflake/cli/plugins/nativeapp/codegen/setup/setup_driver.py.source

@@ -0,0 +1,56 @@
+# 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.
+
+import contextlib
+import os
+import sys
+from pathlib import Path
+
+import snowflake.app.context as ctx
+from snowflake.app.sql import SQLGenerator
+
+ctx._project_path = os.environ["_SNOWFLAKE_CLI_PROJECT_PATH"]
+ctx._current_app_name = os.environ["_SNOWFLAKE_CLI_APP_NAME"]
+__snowflake_internal_py_files = os.environ["_SNOWFLAKE_CLI_SETUP_FILES"].split(
+    os.pathsep
+)
+__snowflake_internal_sql_dest_dir = os.environ["_SNOWFLAKE_CLI_SQL_DEST_DIR"]
+
+try:
+    import importlib
+
+    with contextlib.redirect_stdout(None):
+        with contextlib.redirect_stderr(None):
+            for __snowflake_internal_py_file in __snowflake_internal_py_files:
+                __snowflake_internal_spec = importlib.util.spec_from_file_location(
+                    "<string>", __snowflake_internal_py_file
+                )
+                __snowflake_internal_module = importlib.util.module_from_spec(
+                    __snowflake_internal_spec
+                )
+                __snowflake_internal_spec.loader.exec_module(
+                    __snowflake_internal_module
+                )
+except Exception as exc:  # Catch any error
+    print("An exception occurred while executing file: ", exc, file=sys.stderr)
+    sys.exit(1)
+
+
+import json
+
+output_dir = Path(__snowflake_internal_sql_dest_dir)
+output_dir.mkdir(exist_ok=True, parents=True)
+path_mappings = SQLGenerator(dest_dir=output_dir).generate()
+
+print(json.dumps(path_mappings, default=str))

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

@@ -20,19 +20,21 @@ from pathlib import Path
 from textwrap import dedent
 from typing import Any, Dict, List, Optional, Set
 
-from click import ClickException
 from pydantic import ValidationError
 from snowflake.cli.api.console import cli_console as cc
 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.api.rendering.jinja 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.artifact_processor import (
+    ArtifactProcessor,
+    is_python_file_artifact,
+)
 from snowflake.cli.plugins.nativeapp.codegen.sandbox import (
     ExecutionEnvironmentType,
     SandboxExecutionError,
@@ -167,16 +169,6 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
     ):
         super().__init__(na_project=na_project)
 
-        assert self._na_project.bundle_root.is_absolute()
-        assert self._na_project.deploy_root.is_absolute()
-        assert self._na_project.generated_root.is_absolute()
-        assert self._na_project.project_root.is_absolute()
-
-        if self._na_project.generated_root.exists():
-            raise ClickException(
-                f"Path {self._na_project.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,
@@ -200,7 +192,9 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
 
         collected_output = []
         collected_sql_files: List[Path] = []
-        for py_file, extension_fns in collected_extension_functions_by_path.items():
+        for py_file, extension_fns in sorted(
+            collected_extension_functions_by_path.items()
+        ):
             sql_file = self.generate_new_sql_file_name(
                 py_file=py_file,
             )
@@ -236,9 +230,13 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
             edit_setup_script_with_exec_imm_sql(
                 collected_sql_files=collected_sql_files,
                 deploy_root=bundle_map.deploy_root(),
-                generated_root=self._na_project.generated_root,
+                generated_root=self._generated_root,
             )
 
+    @property
+    def _generated_root(self):
+        return self._na_project.generated_root / "snowpark"
+
     def _normalize_imports(
         self,
         extension_fn: NativeAppExtensionFunction,
@@ -326,8 +324,12 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
             Path, List[NativeAppExtensionFunction]
         ] = {}
 
-        for src_file, dest_file in bundle_map.all_mappings(
-            absolute=True, expand_directories=True, predicate=_is_python_file_artifact
+        for src_file, dest_file in sorted(
+            bundle_map.all_mappings(
+                absolute=True,
+                expand_directories=True,
+                predicate=is_python_file_artifact,
+            )
         ):
             cc.step(
                 "Processing Snowpark annotations from {}".format(
@@ -368,9 +370,7 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
         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._na_project.deploy_root)
-        sql_file = Path(
-            self._na_project.generated_root, relative_py_file.with_suffix(".sql")
-        )
+        sql_file = Path(self._generated_root, relative_py_file.with_suffix(".sql"))
         if sql_file.exists():
             cc.warning(
                 f"""\
@@ -488,7 +488,7 @@ def edit_setup_script_with_exec_imm_sql(
     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 = Path(generated_root, f"__generated.sql")
     generated_file_path.parent.mkdir(exist_ok=True, parents=True)
 
     if generated_file_path.exists():