Add support for --env overrides from command line for templating (#1242)

RELEASE-NOTES.md

@@ -25,6 +25,15 @@
 * Added `--auto-compress` flag to `snow stage copy` command enabling use of gzip to compress files during upload.
 * Added new `native_app.application.post_deploy` section to `snowflake.yml` schema to execute actions after the application has been deployed via `snow app run`.
   * Added the `sql_script` hook type to run SQL scripts with template support.
+* Added support for `--env` command line arguments for templating.
+  * Available for commands that make use of the project definition file.
+  * Format of the argument: `--env key1=value1 --env key2=value2`.
+  * Overrides `env` variables values when used in templating.
+  * Can be referenced in templating through `ctx.env.<key_name>`.
+  * Templating will read env vars in this order of priority (highest priority to lowest priority):
+    * vars from `--env` command line argument.
+    * vars from shell environment variables.
+    * vars from `env` section of project definition file.
 
 ## Fixes and improvements
 * Passing a directory to `snow app deploy` will now deploy any contained file or subfolder specified in the application's artifact rules

src/snowflake/cli/api/cli_global_context.py

@@ -16,7 +16,7 @@
 
 import re
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Callable, Optional
 
 from snowflake.cli.api.exceptions import InvalidSchemaError
 from snowflake.cli.api.output.formats import OutputFormat
@@ -226,6 +226,9 @@ def __init__(self):
         self._experimental = False
         self._project_definition = None
         self._project_root = None
+        self._project_path_arg = None
+        self._project_env_overrides_args = {}
+        self._typer_pre_execute_commands = []
         self._template_context = None
         self._silent: bool = False
 
@@ -261,7 +264,7 @@ def set_experimental(self, value: bool):
         self._experimental = value
 
     @property
-    def project_definition(self) -> Optional[Dict]:
+    def project_definition(self) -> Optional[ProjectDefinition]:
         return self._project_definition
 
     def set_project_definition(self, value: ProjectDefinition):
@@ -274,13 +277,38 @@ def project_root(self):
     def set_project_root(self, project_root: Path):
         self._project_root = project_root
 
+    @property
+    def project_path_arg(self) -> Optional[str]:
+        return self._project_path_arg
+
+    def set_project_path_arg(self, project_path_arg: str):
+        self._project_path_arg = project_path_arg
+
+    @property
+    def project_env_overrides_args(self) -> dict[str, str]:
+        return self._project_env_overrides_args
+
+    def set_project_env_overrides_args(
+        self, project_env_overrides_args: dict[str, str]
+    ):
+        self._project_env_overrides_args = project_env_overrides_args
+
     @property
     def template_context(self) -> dict:
         return self._template_context
 
     def set_template_context(self, template_context: dict):
         self._template_context = template_context
 
+    @property
+    def typer_pre_execute_commands(self) -> list[Callable[[], None]]:
+        return self._typer_pre_execute_commands
+
+    def add_typer_pre_execute_commands(
+        self, typer_pre_execute_command: Callable[[], None]
+    ):
+        self._typer_pre_execute_commands.append(typer_pre_execute_command)
+
     @property
     def connection_context(self) -> _ConnectionContext:
         return self._connection_context

src/snowflake/cli/api/commands/decorators.py

@@ -43,7 +43,8 @@
     VerboseOption,
     WarehouseOption,
     experimental_option,
-    project_type_option,
+    project_definition_option,
+    project_env_overrides_option,
 )
 from snowflake.cli.api.exceptions import CommandReturnTypeError
 from snowflake.cli.api.output.formats import OutputFormat
@@ -72,17 +73,26 @@ def global_options_with_connection(func: Callable):
     )
 
 
-def with_project_definition(project_name: str):
+def with_project_definition(
+    project_name: Optional[str] = None, is_optional: bool = False
+):
     def _decorator(func: Callable):
+
         return _options_decorator_factory(
             func,
             additional_options=[
                 inspect.Parameter(
                     "project_definition",
                     inspect.Parameter.KEYWORD_ONLY,
                     annotation=Optional[str],
-                    default=project_type_option(project_name),
-                )
+                    default=project_definition_option(project_name, is_optional),
+                ),
+                inspect.Parameter(
+                    "env_overrides",
+                    inspect.Parameter.KEYWORD_ONLY,
+                    annotation=List[str],
+                    default=project_env_overrides_option(),
+                ),
             ],
         )
 
@@ -115,7 +125,7 @@ def decorator(func: Callable):
     return decorator
 
 
-def _execute_before_command_using_global_options():
+def _execute_before_command_using_global_options(**options):
     from snowflake.cli.app.loggers import create_loggers
 
     create_loggers(cli_context.verbose, cli_context.enable_tracebacks)
@@ -136,10 +146,15 @@ def _options_decorator_factory(
     additional_options: List[inspect.Parameter],
     execute_before_command_using_new_options: Optional[Callable] = None,
 ):
+    """
+    execute_before_command_using_new_options executes before command telemetry has been emitted,
+    but after command line options have been populated.
+    """
+
     @wraps(func)
     def wrapper(**options):
         if execute_before_command_using_new_options:
-            execute_before_command_using_new_options()
+            execute_before_command_using_new_options(**options)
         return func(**options)
 
     wrapper.__signature__ = _extend_signature_with_additional_options(func, additional_options)  # type: ignore

src/snowflake/cli/api/commands/flags.py

@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import os
 import tempfile
 from dataclasses import dataclass
 from enum import Enum
@@ -26,9 +25,11 @@
 import typer
 from click import ClickException
 from snowflake.cli.api.cli_global_context import cli_context_manager
+from snowflake.cli.api.commands.typer_pre_execute import register_pre_execute_command
 from snowflake.cli.api.console import cli_console
-from snowflake.cli.api.exceptions import MissingConfiguration
+from snowflake.cli.api.exceptions import MissingConfiguration, NoProjectDefinitionError
 from snowflake.cli.api.output.formats import OutputFormat
+from snowflake.cli.api.project.definition_manager import DefinitionManager
 from snowflake.cli.api.utils.rendering import CONTEXT_KEY
 
 DEFAULT_CONTEXT_SETTINGS = {"help_option_names": ["--help", "-h"]}
@@ -499,70 +500,74 @@ def execution_identifier_argument(sf_object: str, example: str) -> typer.Argumen
     )
 
 
-def project_type_option(project_name: str):
-    from snowflake.cli.api.exceptions import NoProjectDefinitionError
-    from snowflake.cli.api.project.definition_manager import DefinitionManager
+def register_project_definition(project_name: Optional[str], is_optional: bool) -> None:
+    project_path = cli_context_manager.project_path_arg
+    env_overrides_args = cli_context_manager.project_env_overrides_args
 
-    def _callback(project_path: Optional[str]):
-        dm = DefinitionManager(project_path)
-        project_definition = dm.project_definition
-        project_root = dm.project_root
+    dm = DefinitionManager(project_path, {CONTEXT_KEY: {"env": env_overrides_args}})
+    project_definition = dm.project_definition
+    project_root = dm.project_root
+    template_context = dm.template_context
 
-        if not getattr(project_definition, project_name, None):
-            raise NoProjectDefinitionError(
-                project_type=project_name, project_file=project_path
-            )
+    if not dm.has_definition_file and not is_optional:
+        raise MissingConfiguration(
+            "Cannot find project definition (snowflake.yml). Please provide a path to the project or run this command in a valid project directory."
+        )
+
+    if project_name is not None and not getattr(project_definition, project_name, None):
+        raise NoProjectDefinitionError(
+            project_type=project_name, project_file=project_path
+        )
+
+    cli_context_manager.set_project_definition(project_definition)
+    cli_context_manager.set_project_root(project_root)
+    cli_context_manager.set_template_context(template_context)
 
-        cli_context_manager.set_project_definition(project_definition)
-        cli_context_manager.set_project_root(project_root)
-        cli_context_manager.set_template_context(dm.template_context)
-        return project_definition
 
-    if project_name == "native_app":
-        project_name_help = "Snowflake Native App"
-    elif project_name == "streamlit":
-        project_name_help = "Streamlit app"
+def _get_project_long_name(project_short_name: Optional[str]) -> str:
+    if project_short_name is None:
+        return "Snowflake"
+
+    if project_short_name == "native_app":
+        project_long_name = "Snowflake Native App"
+    elif project_short_name == "streamlit":
+        project_long_name = "Streamlit app"
     else:
-        project_name_help = project_name.replace("_", " ").capitalize()
+        project_long_name = project_short_name.replace("_", " ").capitalize()
+
+    return f"the {project_long_name}"
+
+
+def project_definition_option(project_name: Optional[str], is_optional: bool):
+    def project_definition_callback(project_path: str) -> None:
+        cli_context_manager.set_project_path_arg(project_path)
+        register_pre_execute_command(
+            lambda: register_project_definition(project_name, is_optional)
+        )
 
     return typer.Option(
         None,
         "-p",
         "--project",
-        help=f"Path where the {project_name_help} project resides. "
+        help=f"Path where {_get_project_long_name(project_name)} project resides. "
         f"Defaults to current working directory.",
-        callback=_callback,
+        callback=_callback(lambda: project_definition_callback),
         show_default=False,
     )
 
 
-def project_definition_option(optional: bool = False):
-    from snowflake.cli.api.project.definition_manager import DefinitionManager
-
-    def _callback(project_path: Optional[str]):
-        try:
-            dm = DefinitionManager(project_path)
-            project_definition = dm.project_definition
-            project_root = dm.project_root
-            template_context = dm.template_context
-        except MissingConfiguration:
-            if optional:
-                project_definition = None
-                project_root = None
-                template_context = {CONTEXT_KEY: {"env": os.environ}}
-            else:
-                raise
-        cli_context_manager.set_project_definition(project_definition)
-        cli_context_manager.set_project_root(project_root)
-        cli_context_manager.set_template_context(template_context)
-        return project_definition
+def project_env_overrides_option():
+    def project_env_overrides_callback(env_overrides_args_list: list[str]) -> None:
+        env_overrides_args_map = {
+            v.key: v.value for v in parse_key_value_variables(env_overrides_args_list)
+        }
+        cli_context_manager.set_project_env_overrides_args(env_overrides_args_map)
 
     return typer.Option(
-        None,
-        "-p",
-        "--project",
-        help=f"Path where Snowflake project resides. Defaults to current working directory.",
-        callback=_callback,
+        [],
+        "--env",
+        help="String in format of key=value. Overrides variables from env section used for templating.",
+        callback=_callback(lambda: project_env_overrides_callback),
         show_default=False,
     )
 
@@ -610,9 +615,13 @@ def __init__(self, key: str, value: str):
         self.value = value
 
 
-def parse_key_value_variables(variables: List[str]) -> List[Variable]:
+def parse_key_value_variables(variables: Optional[List[str]]) -> List[Variable]:
     """Util for parsing key=value input. Useful for commands accepting multiple input options."""
-    result = []
+    result: List[Variable] = []
+
+    if variables is None:
+        return result
+
     for p in variables:
         if "=" not in p:
             raise ClickException(f"Invalid variable: '{p}'")

src/snowflake/cli/api/commands/snow_typer.py

@@ -25,6 +25,7 @@
     global_options_with_connection,
 )
 from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS
+from snowflake.cli.api.commands.typer_pre_execute import run_pre_execute_commands
 from snowflake.cli.api.exceptions import CommandReturnTypeError
 from snowflake.cli.api.output.types import CommandResult
 from snowflake.cli.api.sanitizers import sanitize_for_terminal
@@ -116,6 +117,7 @@ def pre_execute():
         from snowflake.cli.app.telemetry import log_command_usage
 
         log.debug("Executing command pre execution callback")
+        run_pre_execute_commands()
         log_command_usage()
 
     @staticmethod

src/snowflake/cli/api/commands/typer_pre_execute.py

@@ -0,0 +1,26 @@
+# 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 typing import Callable
+
+from snowflake.cli.api.cli_global_context import cli_context_manager
+
+
+def register_pre_execute_command(command: Callable[[], None]) -> None:
+    cli_context_manager.add_typer_pre_execute_commands(command)
+
+
+def run_pre_execute_commands() -> None:
+    for command in cli_context_manager.typer_pre_execute_commands:
+        command()