snowflake-cli-labs 3.0.0rc0__py3-none-any.whl → 3.0.0rc1__py3-none-any.whl

snowflake/cli/__about__.py

@@ -14,4 +14,4 @@
 
 from __future__ import annotations
 
-VERSION = "3.0.0rc0"
+VERSION = "3.0.0rc1"

snowflake/cli/_app/snow_connector.py

@@ -71,6 +71,7 @@ def connect_to_snowflake(
 
     if connection_name:
         connection_parameters = get_connection_dict(connection_name)
+        connection_parameters = get_connection_dict(connection_name)
     elif temporary_connection:
         connection_parameters = {}  # we will apply overrides in next step
     else:
@@ -164,18 +165,24 @@ def _raise_errors_related_to_session_token(
 
 
 def update_connection_details_with_private_key(connection_parameters: Dict):
-    if "private_key_path" in connection_parameters:
-        if connection_parameters.get("authenticator") == "SNOWFLAKE_JWT":
-            private_key = _load_pem_to_der(connection_parameters["private_key_path"])
-            connection_parameters["private_key"] = private_key
-            del connection_parameters["private_key_path"]
-        else:
-            raise ClickException(
-                "Private Key authentication requires authenticator set to SNOWFLAKE_JWT"
-            )
+    if "private_key_file" in connection_parameters:
+        _load_private_key(connection_parameters, "private_key_file")
+    elif "private_key_path" in connection_parameters:
+        _load_private_key(connection_parameters, "private_key_path")
     return connection_parameters
 
 
+def _load_private_key(connection_parameters: Dict, private_key_var_name: str) -> None:
+    if connection_parameters.get("authenticator") == "SNOWFLAKE_JWT":
+        private_key = _load_pem_to_der(connection_parameters[private_key_var_name])
+        connection_parameters["private_key"] = private_key
+        del connection_parameters[private_key_var_name]
+    else:
+        raise ClickException(
+            "Private Key authentication requires authenticator set to SNOWFLAKE_JWT"
+        )
+
+
 def _update_connection_application_name(connection_parameters: Dict):
     """Update version and name of app handling connection."""
     connection_application_params = {
@@ -184,13 +191,13 @@ def _update_connection_application_name(connection_parameters: Dict):
     connection_parameters.update(connection_application_params)
 
 
-def _load_pem_to_der(private_key_path: str) -> bytes:
+def _load_pem_to_der(private_key_file: str) -> bytes:
     """
     Given a private key file path (in PEM format), decode key data into DER
     format
     """
 
-    with SecurePath(private_key_path).open(
+    with SecurePath(private_key_file).open(
         "rb", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB
     ) as f:
         private_key_pem = f.read()

snowflake/cli/_plugins/connection/commands.py

@@ -224,9 +224,10 @@ def add(
         prompt="Authentication method",
         help="Chosen authenticator, if other than password-based",
     ),
-    private_key_path: str = typer.Option(
+    private_key_file: str = typer.Option(
         EmptyInput(),
         "--private-key",
+        "--private-key-path",
         "-k",
         click_type=OptionalPrompt(),
         prompt="Path to private key file",
@@ -268,7 +269,7 @@ def add(
             warehouse=warehouse,
             role=role,
             authenticator=authenticator,
-            private_key_path=private_key_path,
+            private_key_file=private_key_file,
             token_file_path=token_file_path,
         ),
     )

snowflake/cli/_plugins/git/manager.py

@@ -41,17 +41,25 @@ class GitStagePathParts(StagePathParts):
 
     @property
     def path(self) -> str:
-        return (
-            f"{self.stage_name}{self.directory}"
-            if self.stage_name.endswith("/")
-            else f"{self.stage_name}/{self.directory}"
-        )
+        return f"{self.stage_name.rstrip('/')}/{self.directory}"
 
-    def add_stage_prefix(self, file_path: str) -> str:
+    @classmethod
+    def get_directory(cls, stage_path: str) -> str:
+        return "/".join(Path(stage_path).parts[3:])
+
+    @property
+    def full_path(self) -> str:
+        return f"{self.stage.rstrip('/')}/{self.directory}"
+
+    def replace_stage_prefix(self, file_path: str) -> str:
         stage = Path(self.stage).parts[0]
         file_path_without_prefix = Path(file_path).parts[1:]
         return f"{stage}/{'/'.join(file_path_without_prefix)}"
 
+    def add_stage_prefix(self, file_path: str) -> str:
+        stage = self.stage.rstrip("/")
+        return f"{stage}/{file_path.lstrip('/')}"
+
     def get_directory_from_file_path(self, file_path: str) -> List[str]:
         stage_path_length = len(Path(self.directory).parts)
         return list(Path(file_path).parts[3 + stage_path_length : -1])

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

@@ -14,8 +14,11 @@
 
 from __future__ import annotations
 
+import copy
+import re
 from typing import Dict, Optional
 
+from click import ClickException
 from snowflake.cli._plugins.nativeapp.bundle_context import BundleContext
 from snowflake.cli._plugins.nativeapp.codegen.artifact_processor import (
     ArtifactProcessor,
@@ -27,6 +30,9 @@ from snowflake.cli._plugins.nativeapp.codegen.setup.native_app_setup_processor i
 from snowflake.cli._plugins.nativeapp.codegen.snowpark.python_processor import (
     SnowparkAnnotationProcessor,
 )
+from snowflake.cli._plugins.nativeapp.codegen.templates.templates_processor import (
+    TemplatesProcessor,
+)
 from snowflake.cli._plugins.nativeapp.feature_flags import FeatureFlag
 from snowflake.cli.api.console import cli_console as cc
 from snowflake.cli.api.project.schemas.native_app.path_mapping import (
@@ -34,11 +40,13 @@ from snowflake.cli.api.project.schemas.native_app.path_mapping import (
 )
 
 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,
 }
 
 
@@ -110,7 +118,15 @@ class NativeAppCompiler:
             # No registered processor with the specified name
             return None
 
-        current_processor = processor_factory(self._bundle_ctx)
+        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

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

@@ -15,12 +15,18 @@
 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_setup_script_file
+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,
@@ -40,6 +46,32 @@ from snowflake.cli.api.project.schemas.native_app.path_mapping import (
 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):
@@ -62,7 +94,7 @@ class NativeAppSetupProcessor(ArtifactProcessor):
 
         self._create_or_update_sandbox()
 
-        cc.phase("Processing Python setup files")
+        cc.step("Processing Python setup files")
 
         files_to_process = []
         for src_file, dest_file in bundle_map.all_mappings(
@@ -73,18 +105,55 @@ class NativeAppSetupProcessor(ArtifactProcessor):
             )
             files_to_process.append(src_file)
 
-        sql_files_mapping = self._execute_in_sandbox(files_to_process)
-        self._generate_setup_sql(sql_files_mapping)
+        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(self.generated_root),
+            "_SNOWFLAKE_CLI_SQL_DEST_DIR": str(generated_root),
+            "_SNOWFLAKE_CLI_MANIFEST_PATH": str(manifest_path),
         }
 
         try:
@@ -102,56 +171,68 @@ class NativeAppSetupProcessor(ArtifactProcessor):
             )
 
         if result.returncode == 0:
-            sql_file_mappings = json.loads(result.stdout)
-            return sql_file_mappings
+            return json.loads(result.stdout)
         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)
-
+    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:
-            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._bundle_ctx.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._bundle_ctx.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._bundle_ctx.deploy_root))}';"
-                )
-
-            f.edited_contents = "\n".join(new_contents)
+        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 / "setup_py_venv"
-
-    @property
-    def generated_root(self):
-        return self._bundle_ctx.generated_root / "setup_py"
+        return self._bundle_ctx.bundle_root / "venv"
 
     def _create_or_update_sandbox(self):
         sandbox_root = self.sandbox_root

snowflake/cli/_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/nativeapp/codegen/snowpark/python_processor.py

@@ -226,13 +226,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._bundle_ctx.generated_root / "snowpark"
-
     def _normalize_imports(
         self,
         extension_fn: NativeAppExtensionFunction,
@@ -366,7 +362,9 @@ 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._bundle_ctx.deploy_root)
-        sql_file = Path(self._generated_root, relative_py_file.with_suffix(".sql"))
+        sql_file = Path(
+            self._bundle_ctx.generated_root, relative_py_file.with_suffix(".sql")
+        )
         if sql_file.exists():
             cc.warning(
                 f"""\

snowflake/cli/_plugins/nativeapp/codegen/templates/templates_processor.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Optional
+
+import jinja2
+from snowflake.cli._plugins.nativeapp.artifacts import BundleMap
+from snowflake.cli._plugins.nativeapp.codegen.artifact_processor import (
+    ArtifactProcessor,
+)
+from snowflake.cli._plugins.nativeapp.exceptions import InvalidTemplateInFileError
+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.project.schemas.native_app.path_mapping import (
+    PathMapping,
+    ProcessorMapping,
+)
+from snowflake.cli.api.rendering.project_definition_templates import (
+    get_client_side_jinja_env,
+)
+from snowflake.cli.api.rendering.sql_templates import (
+    choose_sql_jinja_env_based_on_template_syntax,
+)
+
+
+class TemplatesProcessor(ArtifactProcessor):
+    """
+    Processor class to perform template expansion on all relevant artifacts (specified in the project definition file).
+    """
+
+    def process(
+        self,
+        artifact_to_process: PathMapping,
+        processor_mapping: Optional[ProcessorMapping],
+        **kwargs,
+    ):
+        """
+        Process the artifact by executing the template expansion logic on it.
+        """
+        cc.step(f"Processing artifact {artifact_to_process} with templates processor")
+
+        bundle_map = BundleMap(
+            project_root=self._bundle_ctx.project_root,
+            deploy_root=self._bundle_ctx.deploy_root,
+        )
+        bundle_map.add(artifact_to_process)
+
+        for src, dest in bundle_map.all_mappings(
+            absolute=True,
+            expand_directories=True,
+        ):
+            if src.is_dir():
+                continue
+            with self.edit_file(dest) as f:
+                file_name = src.relative_to(self._bundle_ctx.project_root)
+
+                jinja_env = (
+                    choose_sql_jinja_env_based_on_template_syntax(
+                        f.contents, reference_name=file_name
+                    )
+                    if dest.name.lower().endswith(".sql")
+                    else get_client_side_jinja_env()
+                )
+
+                try:
+                    expanded_template = jinja_env.from_string(f.contents).render(
+                        get_cli_context().template_context
+                    )
+
+                # For now, we are printing the source file path in the error message
+                # instead of the destination file path to make it easier for the user
+                # to identify the file that has the error, and edit the correct file.
+                except jinja2.TemplateSyntaxError as e:
+                    raise InvalidTemplateInFileError(file_name, e, e.lineno) from e
+
+                except jinja2.UndefinedError as e:
+                    raise InvalidTemplateInFileError(file_name, e) from e
+
+                if expanded_template != f.contents:
+                    f.edited_contents = expanded_template

snowflake/cli/_plugins/nativeapp/exceptions.py

@@ -64,15 +64,15 @@ class MissingScriptError(ClickException):
         super().__init__(f'Script "{relpath}" does not exist')
 
 
-class InvalidScriptError(ClickException):
-    """A referenced script had syntax error(s)."""
+class InvalidTemplateInFileError(ClickException):
+    """A referenced templated file had syntax error(s)."""
 
     def __init__(
         self, relpath: str, err: jinja2.TemplateError, lineno: Optional[int] = None
     ):
         lineno_str = f":{lineno}" if lineno is not None else ""
         super().__init__(
-            f'Script "{relpath}{lineno_str}" does not contain a valid template: {err.message}'
+            f'File "{relpath}{lineno_str}" does not contain a valid template: {err.message}'
         )
         self.err = err