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

snowflake/cli/_plugins/nativeapp/bundle_context.py

@@ -0,0 +1,31 @@
+# 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 dataclasses import dataclass
+from pathlib import Path
+from typing import (
+    List,
+)
+
+from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import PathMapping
+
+
+@dataclass
+class BundleContext:
+    package_name: str
+    artifacts: List[PathMapping]
+    project_root: Path
+    bundle_root: Path
+    deploy_root: Path
+    generated_root: Path

snowflake/cli/{plugins → _plugins}/nativeapp/codegen/artifact_processor.py

@@ -19,11 +19,11 @@ from pathlib import Path
 from typing import Optional
 
 from click import ClickException
-from snowflake.cli.api.project.schemas.native_app.path_mapping import (
+from snowflake.cli._plugins.nativeapp.bundle_context import BundleContext
+from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import (
     PathMapping,
     ProcessorMapping,
 )
-from snowflake.cli.plugins.nativeapp.project_model import NativeAppProjectModel
 
 
 class UnsupportedArtifactProcessorError(ClickException):
@@ -74,9 +74,9 @@ class ProjectFileContextManager:
 class ArtifactProcessor(ABC):
     def __init__(
         self,
-        na_project: NativeAppProjectModel,
+        bundle_ctx: BundleContext,
     ) -> None:
-        self._na_project = na_project
+        self._bundle_ctx = bundle_ctx
 
     @abstractmethod
     def process(

snowflake/cli/{plugins → _plugins}/nativeapp/codegen/compiler.py

@@ -14,31 +14,41 @@
 
 from __future__ import annotations
 
+import copy
+import re
 from typing import Dict, Optional
 
-from snowflake.cli.api.console import cli_console as cc
-from snowflake.cli.api.project.schemas.native_app.path_mapping import (
-    ProcessorMapping,
-)
-from snowflake.cli.plugins.nativeapp.codegen.artifact_processor import (
+from click import ClickException
+from snowflake.cli._plugins.nativeapp.bundle_context import BundleContext
+from snowflake.cli._plugins.nativeapp.codegen.artifact_processor import (
     ArtifactProcessor,
     UnsupportedArtifactProcessorError,
 )
-from snowflake.cli.plugins.nativeapp.codegen.setup.native_app_setup_processor import (
+from snowflake.cli._plugins.nativeapp.codegen.setup.native_app_setup_processor import (
     NativeAppSetupProcessor,
 )
-from snowflake.cli.plugins.nativeapp.codegen.snowpark.python_processor import (
+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
+from snowflake.cli._plugins.nativeapp.codegen.templates.templates_processor import (
+    TemplatesProcessor,
+)
+from snowflake.cli._plugins.nativeapp.feature_flags import FeatureFlag
+from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.console import cli_console as cc
+from snowflake.cli.api.metrics import CLICounterField
+from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import (
+    ProcessorMapping,
+)
 
 SNOWPARK_PROCESSOR = "snowpark"
-NA_SETUP_PROCESSOR = "native-app-setup"
+NA_SETUP_PROCESSOR = "native app setup"
+TEMPLATES_PROCESSOR = "templates"
 
 _REGISTERED_PROCESSORS_BY_NAME = {
     SNOWPARK_PROCESSOR: SnowparkAnnotationProcessor,
     NA_SETUP_PROCESSOR: NativeAppSetupProcessor,
+    TEMPLATES_PROCESSOR: TemplatesProcessor,
 }
 
 
@@ -53,9 +63,9 @@ class NativeAppCompiler:
 
     def __init__(
         self,
-        na_project: NativeAppProjectModel,
+        bundle_ctx: BundleContext,
     ):
-        self._na_project = na_project
+        self._bundle_ctx = bundle_ctx
         # dictionary of all processors created and shared between different artifact objects.
         self.cached_processors: Dict[str, ArtifactProcessor] = {}
 
@@ -64,17 +74,20 @@ 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.
         """
+        metrics = get_cli_context().metrics
+        metrics.set_counter_default(CLICounterField.TEMPLATES_PROCESSOR, 0)
+        metrics.set_counter_default(CLICounterField.SNOWPARK_PROCESSOR, 0)
 
         if not self._should_invoke_processors():
             return
 
         with cc.phase("Invoking artifact processors"):
-            if self._na_project.generated_root.exists():
+            if self._bundle_ctx.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."
+                    f"Path {self._bundle_ctx.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 artifact in self._bundle_ctx.artifacts:
                 for processor in artifact.processors:
                     if self._is_enabled(processor):
                         artifact_processor = self._try_create_processor(
@@ -110,15 +123,21 @@ class NativeAppCompiler:
             # No registered processor with the specified name
             return None
 
-        current_processor = processor_factory(
-            na_project=self._na_project,
+        processor_ctx = copy.copy(self._bundle_ctx)
+        processor_subdirectory = re.sub(r"[^a-zA-Z0-9_$]", "_", processor_name)
+        processor_ctx.bundle_root = (
+            self._bundle_ctx.bundle_root / processor_subdirectory
+        )
+        processor_ctx.generated_root = (
+            self._bundle_ctx.generated_root / processor_subdirectory
         )
+        current_processor = processor_factory(processor_ctx)
         self.cached_processors[processor_name] = current_processor
 
         return current_processor
 
     def _should_invoke_processors(self):
-        for artifact in self._na_project.artifacts:
+        for artifact in self._bundle_ctx.artifacts:
             for processor in artifact.processors:
                 if self._is_enabled(processor):
                     return True

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

@@ -0,0 +1,249 @@
+# 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 logging
+import os.path
+from pathlib import Path
+from typing import List, Optional
+
+import yaml
+from click import ClickException
+from snowflake.cli._plugins.nativeapp.artifacts import (
+    BundleMap,
+    find_manifest_file,
+    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.stage.diff import to_stage_path
+from snowflake.cli.api.console import cli_console as cc
+from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import (
+    PathMapping,
+    ProcessorMapping,
+)
+
+DEFAULT_TIMEOUT = 30
+DRIVER_PATH = Path(__file__).parent / "setup_driver.py.source"
+
+log = logging.getLogger(__name__)
+
+
+def safe_set(d: dict, *keys: str, **kwargs) -> None:
+    """
+    Sets a value in a nested dictionary structure, creating intermediate dictionaries as needed.
+    Sample usage:
+
+      d = {}
+      safe_set(d, "a", "b", "c", value=42)
+
+    d is now:
+      {
+        "a": {
+          "b": {
+            "c": 42
+          }
+        }
+      }
+    """
+    curr = d
+    for k in keys[:-1]:
+        curr = curr.setdefault(k, {})
+
+    curr[keys[-1]] = kwargs.get("value")
+
+
+class NativeAppSetupProcessor(ArtifactProcessor):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    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._bundle_ctx.project_root,
+            deploy_root=self._bundle_ctx.deploy_root,
+        )
+        bundle_map.add(artifact_to_process)
+
+        self._create_or_update_sandbox()
+
+        cc.step("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._bundle_ctx.project_root)}"
+            )
+            files_to_process.append(src_file)
+
+        result = self._execute_in_sandbox(files_to_process)
+        if not result:
+            return  # nothing to do
+
+        logs = result.get("logs", [])
+        for msg in logs:
+            log.debug(msg)
+
+        warnings = result.get("warnings", [])
+        for msg in warnings:
+            cc.warning(msg)
+
+        schema_version = result.get("schema_version")
+        if schema_version != "1":
+            raise ClickException(
+                f"Unsupported schema version returned from snowflake-app-python library: {schema_version}"
+            )
+
+        setup_script_mods = [
+            mod
+            for mod in result.get("modifications", [])
+            if mod.get("target") == "native_app:setup_script"
+        ]
+        if setup_script_mods:
+            self._edit_setup_sql(setup_script_mods)
+
+        manifest_mods = [
+            mod
+            for mod in result.get("modifications", [])
+            if mod.get("target") == "native_app:manifest"
+        ]
+        if manifest_mods:
+            self._edit_manifest(manifest_mods)
+
+    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 ''}")
+
+        manifest_path = find_manifest_file(deploy_root=self._bundle_ctx.deploy_root)
+
+        generated_root = self._bundle_ctx.generated_root
+        generated_root.mkdir(exist_ok=True, parents=True)
+
+        env_vars = {
+            "_SNOWFLAKE_CLI_PROJECT_PATH": str(self._bundle_ctx.project_root),
+            "_SNOWFLAKE_CLI_SETUP_FILES": os.pathsep.join(map(str, py_files)),
+            "_SNOWFLAKE_CLI_APP_NAME": str(self._bundle_ctx.package_name),
+            "_SNOWFLAKE_CLI_SQL_DEST_DIR": str(generated_root),
+            "_SNOWFLAKE_CLI_MANIFEST_PATH": str(manifest_path),
+        }
+
+        try:
+            result = execute_script_in_sandbox(
+                script_source=DRIVER_PATH.read_text(),
+                env_type=ExecutionEnvironmentType.VENV,
+                cwd=self._bundle_ctx.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:
+            return json.loads(result.stdout)
+        else:
+            raise ClickException(
+                f"Failed to execute python setup script logic: {result.stderr}"
+            )
+
+    def _edit_setup_sql(self, modifications: List[dict]) -> None:
+        cc.step("Patching setup script")
+        setup_file_path = find_setup_script_file(
+            deploy_root=self._bundle_ctx.deploy_root
+        )
+
+        with self.edit_file(setup_file_path) as f:
+            prepended = []
+            appended = []
+
+            for mod in modifications:
+                for inst in mod.get("instructions", []):
+                    if inst.get("type") == "insert":
+                        default_loc = inst.get("default_location")
+                        if default_loc == "end":
+                            appended.append(self._setup_mod_instruction_to_sql(inst))
+                        elif default_loc == "start":
+                            prepended.append(self._setup_mod_instruction_to_sql(inst))
+
+            if prepended or appended:
+                f.edited_contents = "\n".join(prepended + [f.contents] + appended)
+
+    def _edit_manifest(self, modifications: List[dict]) -> None:
+        cc.step("Patching manifest")
+        manifest_path = find_manifest_file(deploy_root=self._bundle_ctx.deploy_root)
+
+        with self.edit_file(manifest_path) as f:
+            manifest = yaml.safe_load(f.contents)
+
+            for mod in modifications:
+                for inst in mod.get("instructions", []):
+                    if inst.get("type") == "set":
+                        payload = inst.get("payload")
+                        if payload:
+                            key = payload.get("key")
+                            value = payload.get("value")
+                            safe_set(manifest, *key.split("."), value=value)
+            f.edited_contents = yaml.safe_dump(manifest, sort_keys=False)
+
+    def _setup_mod_instruction_to_sql(self, mod_inst: dict) -> str:
+        payload = mod_inst.get("payload")
+        if not payload:
+            raise ClickException("Unsupported instruction received: no payload found")
+
+        payload_type = payload.get("type")
+        if payload_type == "execute immediate":
+            file_path = payload.get("file_path")
+            if file_path:
+                sql_file_path = self._bundle_ctx.generated_root / file_path
+                return f"EXECUTE IMMEDIATE FROM '/{to_stage_path(sql_file_path.relative_to(self._bundle_ctx.deploy_root))}';"
+
+        raise ClickException(f"Unsupported instruction type received: {payload_type}")
+
+    @property
+    def sandbox_root(self):
+        return self._bundle_ctx.bundle_root / "venv"
+
+    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._bundle_ctx.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 → _plugins}/nativeapp/codegen/setup/setup_driver.py.source

@@ -20,8 +20,11 @@ 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"]
+ctx.configure("project_path", os.environ.get("_SNOWFLAKE_CLI_PROJECT_PATH", None))
+ctx.configure("manifest_path", os.environ.get("_SNOWFLAKE_CLI_MANIFEST_PATH", None))
+ctx.configure("current_app_name", os.environ.get("_SNOWFLAKE_CLI_APP_NAME", None))
+ctx.configure("enable_sql_generation", True)
+
 __snowflake_internal_py_files = os.environ["_SNOWFLAKE_CLI_SETUP_FILES"].split(
     os.pathsep
 )

snowflake/cli/{plugins → _plugins}/nativeapp/codegen/snowpark/extension_function_utils.py

@@ -23,17 +23,17 @@ from typing import (
 )
 
 from click.exceptions import ClickException
-from snowflake.cli.api.project.schemas.snowpark.argument import Argument
+from snowflake.cli._plugins.nativeapp.codegen.snowpark.models import (
+    ExtensionFunctionTypeEnum,
+    NativeAppExtensionFunction,
+)
+from snowflake.cli.api.project.schemas.v1.snowpark.argument import Argument
 from snowflake.cli.api.project.util import (
     is_valid_identifier,
     is_valid_string_literal,
     to_identifier,
     to_string_literal,
 )
-from snowflake.cli.plugins.nativeapp.codegen.snowpark.models import (
-    ExtensionFunctionTypeEnum,
-    NativeAppExtensionFunction,
-)
 
 ASTDefNode = Union[ast.FunctionDef, ast.ClassDef]
 

snowflake/cli/{plugins → _plugins}/nativeapp/codegen/snowpark/models.py

@@ -18,8 +18,8 @@ from enum import Enum
 from typing import List, Optional
 
 from pydantic import Field
-from snowflake.cli.api.project.schemas.snowpark.callable import _CallableBase
 from snowflake.cli.api.project.schemas.updatable_model import IdentifierField
+from snowflake.cli.api.project.schemas.v1.snowpark.callable import _CallableBase
 
 
 class ExtensionFunctionTypeEnum(str, Enum):

snowflake/cli/{plugins → _plugins}/nativeapp/codegen/snowpark/python_processor.py

@@ -21,26 +21,20 @@ from textwrap import dedent
 from typing import Any, Dict, List, Optional, Set
 
 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.rendering.jinja import jinja_render_from_file
-from snowflake.cli.plugins.nativeapp.artifacts import (
+from snowflake.cli._plugins.nativeapp.artifacts import (
     BundleMap,
     find_setup_script_file,
 )
-from snowflake.cli.plugins.nativeapp.codegen.artifact_processor import (
+from snowflake.cli._plugins.nativeapp.codegen.artifact_processor import (
     ArtifactProcessor,
     is_python_file_artifact,
 )
-from snowflake.cli.plugins.nativeapp.codegen.sandbox import (
+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 (
+from snowflake.cli._plugins.nativeapp.codegen.snowpark.extension_function_utils import (
     deannotate_module_source,
     ensure_all_string_literals,
     ensure_string_literal,
@@ -49,12 +43,19 @@ from snowflake.cli.plugins.nativeapp.codegen.snowpark.extension_function_utils i
     get_sql_argument_signature,
     get_sql_object_type,
 )
-from snowflake.cli.plugins.nativeapp.codegen.snowpark.models import (
+from snowflake.cli._plugins.nativeapp.codegen.snowpark.models import (
     ExtensionFunctionTypeEnum,
     NativeAppExtensionFunction,
 )
-from snowflake.cli.plugins.nativeapp.project_model import NativeAppProjectModel
-from snowflake.cli.plugins.stage.diff import to_stage_path
+from snowflake.cli._plugins.stage.diff import to_stage_path
+from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.console import cli_console as cc
+from snowflake.cli.api.metrics import CLICounterField
+from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import (
+    PathMapping,
+    ProcessorMapping,
+)
+from snowflake.cli.api.rendering.jinja import jinja_render_from_file
 
 DEFAULT_TIMEOUT = 30
 TEMPLATE_PATH = Path(__file__).parent / "callback_source.py.jinja"
@@ -163,11 +164,8 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
     and generate SQL code for creation of extension functions based on those discovered objects.
     """
 
-    def __init__(
-        self,
-        na_project: NativeAppProjectModel,
-    ):
-        super().__init__(na_project=na_project)
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
 
     def process(
         self,
@@ -180,9 +178,11 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
         setup script with generated SQL that registers these functions.
         """
 
+        get_cli_context().metrics.set_counter(CLICounterField.SNOWPARK_PROCESSOR, 1)
+
         bundle_map = BundleMap(
-            project_root=self._na_project.project_root,
-            deploy_root=self._na_project.deploy_root,
+            project_root=self._bundle_ctx.project_root,
+            deploy_root=self._bundle_ctx.deploy_root,
         )
         bundle_map.add(artifact_to_process)
 
@@ -230,13 +230,9 @@ 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._generated_root,
+                generated_root=self._bundle_ctx.generated_root,
             )
 
-    @property
-    def _generated_root(self):
-        return self._na_project.generated_root / "snowpark"
-
     def _normalize_imports(
         self,
         extension_fn: NativeAppExtensionFunction,
@@ -315,7 +311,7 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
         self, bundle_map: BundleMap, processor_mapping: Optional[ProcessorMapping]
     ) -> Dict[Path, List[NativeAppExtensionFunction]]:
         kwargs = (
-            _determine_virtual_env(self._na_project.project_root, processor_mapping)
+            _determine_virtual_env(self._bundle_ctx.project_root, processor_mapping)
             if processor_mapping is not None
             else {}
         )
@@ -331,14 +327,11 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
                 predicate=is_python_file_artifact,
             )
         ):
-            cc.step(
-                "Processing Snowpark annotations from {}".format(
-                    dest_file.relative_to(bundle_map.deploy_root())
-                )
-            )
+            src_file_name = src_file.relative_to(self._bundle_ctx.project_root)
+            cc.step(f"Processing Snowpark annotations from {src_file_name}")
             collected_extension_function_json = _execute_in_sandbox(
                 py_file=str(dest_file.resolve()),
-                deploy_root=self._na_project.deploy_root,
+                deploy_root=self._bundle_ctx.deploy_root,
                 kwargs=kwargs,
             )
 
@@ -369,8 +362,10 @@ 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._generated_root, relative_py_file.with_suffix(".sql"))
+        relative_py_file = py_file.relative_to(self._bundle_ctx.deploy_root)
+        sql_file = Path(
+            self._bundle_ctx.generated_root, relative_py_file.with_suffix(".sql")
+        )
         if sql_file.exists():
             cc.warning(
                 f"""\