Allow use of environment variables in deployment steps (#10199)

Co-authored-by: Chris Pickett <[email protected]>

docs/concepts/deployments-ux.md

@@ -295,11 +295,12 @@ You can also run custom steps by packaging them. In the example below, `retrieve
 
 ### Templating Options
 
-Values that you place within your `prefect.yaml` file can reference dynamic values in two different ways:
+Values that you place within your `prefect.yaml` file can reference dynamic values in several different ways:
 
 - **step outputs**: every step of both `build` and `push` produce named fields such as `image_name`; you can reference these fields within `prefect.yaml` and `prefect deploy` will populate them with each call.  References must be enclosed in double brackets and be of the form `"{{ field_name }}"`
 - **blocks**: [Prefect blocks](/concepts/blocks) can also be referenced with the special syntax `{{ prefect.blocks.block_type.block_slug }}`; it is highly recommended that you use block references for any sensitive information (such as a GitHub access token or any credentials) to avoid hardcoding these values in plaintext
 - **variables**: [Prefect variables](/concepts/variables) can also be referenced with the special syntax `{{ prefect.variables.variable_name }}`. Variables can be used to reference non-sensitive, reusable pieces of information such as a default image name or a default work pool name.
+- **environment variables**: you can also reference environment variables with the special syntax `{{ $MY_ENV_VAR }}`. This is especially useful for referencing environment variables that are set at runtime.
 
 As an example, consider the following `prefect.yaml` file:
 
@@ -316,9 +317,9 @@ build:
 deployments:
   - # base metadata
     name: null
-    version: "{{ build_image.tag }}"
+    version: "{{ build-image.tag }}"
     tags:
-        - "{{ build_image.tag }}"
+        - "{{ $my_deployment_tag }}"
         - "{{ prefect.variables.some_common_tag }}"
     description: null
     schedule: null
@@ -333,7 +334,7 @@ deployments:
         name: "my-k8s-work-pool"
         work_queue_name: null
         job_variables:
-            image: "{{ build_image.image }}"
+            image: "{{ build-image.image }}"
             cluster_config: "{{ prefect.blocks.kubernetes-cluster-config.my-favorite-config }}"
 ```
 

src/prefect/deployments/steps/core.py

@@ -11,6 +11,7 @@
 - The step's output is returned and used to resolve inputs for subsequent steps
 """
 from copy import deepcopy
+import os
 import subprocess
 import sys
 from typing import Any, Dict, List, Optional, Tuple
@@ -89,6 +90,7 @@ async def run_step(step: Dict, upstream_outputs: Optional[Dict] = None) -> Dict:
     inputs = apply_values(inputs, upstream_outputs)
     inputs = await resolve_block_document_references(inputs)
     inputs = await resolve_variables(inputs)
+    inputs = apply_values(inputs, os.environ)
     step_func = _get_function_for_step(fqn, requires=keywords.get("requires"))
     result = await from_async.call_soon_in_new_thread(
         Call.new(step_func, **inputs)

src/prefect/utilities/templating.py

@@ -1,4 +1,5 @@
 import enum
+import os
 import re
 from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Set, Type, TypeVar, Union
 
@@ -13,15 +14,17 @@
 
 T = TypeVar("T", str, int, float, bool, dict, list, None)
 
-PLACEHOLDER_CAPTURE_REGEX = re.compile(r"({{\s*([\w\.\-\[\]]+)\s*}})")
+PLACEHOLDER_CAPTURE_REGEX = re.compile(r"({{\s*([\w\.\-\[\]$]+)\s*}})")
 BLOCK_DOCUMENT_PLACEHOLDER_PREFIX = "prefect.blocks."
 VARIABLE_PLACEHOLDER_PREFIX = "prefect.variables."
+ENV_VAR_PLACEHOLDER_PREFIX = "$"
 
 
 class PlaceholderType(enum.Enum):
     STANDARD = "standard"
     BLOCK_DOCUMENT = "block_document"
     VARIABLE = "variable"
+    ENV_VAR = "env_var"
 
 
 class Placeholder(NamedTuple):
@@ -44,6 +47,8 @@ def determine_placeholder_type(name: str) -> PlaceholderType:
         return PlaceholderType.BLOCK_DOCUMENT
     elif name.startswith(VARIABLE_PLACEHOLDER_PREFIX):
         return PlaceholderType.VARIABLE
+    elif name.startswith(ENV_VAR_PLACEHOLDER_PREFIX):
+        return PlaceholderType.ENV_VAR
     else:
         return PlaceholderType.STANDARD
 
@@ -126,12 +131,19 @@ def apply_values(
             for full_match, name, placeholder_type in placeholders:
                 if placeholder_type is PlaceholderType.STANDARD:
                     value = get_from_dict(values, name, NotSet)
-                    if value is NotSet and not remove_notset:
-                        continue
-                    elif value is NotSet:
-                        template = template.replace(full_match, "")
-                    else:
-                        template = template.replace(full_match, str(value))
+                elif placeholder_type is PlaceholderType.ENV_VAR:
+                    name = name.lstrip(ENV_VAR_PLACEHOLDER_PREFIX)
+                    value = os.environ.get(name, NotSet)
+                else:
+                    continue
+
+                if value is NotSet and not remove_notset:
+                    continue
+                elif value is NotSet:
+                    template = template.replace(full_match, "")
+                else:
+                    template = template.replace(full_match, str(value))
+
             return template
     elif isinstance(template, dict):
         updated_template = {}

tests/cli/test_deploy.py

@@ -1307,6 +1307,68 @@ async def test_project_deploy_templates_values(self, work_pool, prefect_client):
         assert deployment.tags == ["b", "2", "3"]
         assert deployment.description == "1"
 
+    @pytest.mark.usefixtures("project_dir")
+    async def test_project_deploy_templates_env_var_values(
+        self, prefect_client, work_pool, monkeypatch
+    ):
+        # prepare a templated deployment
+        prefect_file = Path("prefect.yaml")
+        with prefect_file.open(mode="r") as f:
+            contents = yaml.safe_load(f)
+
+        contents["deployments"][0]["name"] = "test-name"
+        contents["deployments"][0]["version"] = "{{ $MY_VERSION }}"
+        contents["deployments"][0]["tags"] = "{{ $MY_TAGS }}"
+        contents["deployments"][0]["description"] = "{{ $MY_DESCRIPTION }}"
+
+        # save it back
+        with prefect_file.open(mode="w") as f:
+            yaml.safe_dump(contents, f)
+
+        # update prefect.yaml to include some new build steps
+        prefect_file = Path("prefect.yaml")
+        with prefect_file.open(mode="r") as f:
+            prefect_config = yaml.safe_load(f)
+
+        monkeypatch.setenv("MY_DIRECTORY", "bar")
+        monkeypatch.setenv("MY_FILE", "foo.txt")
+
+        prefect_config["build"] = [
+            {
+                "prefect.deployments.steps.run_shell_script": {
+                    "id": "get-dir",
+                    "script": "echo '{{ $MY_DIRECTORY }}'",
+                    "stream_output": True,
+                }
+            },
+        ]
+
+        # save it back
+        with prefect_file.open(mode="w") as f:
+            yaml.safe_dump(prefect_config, f)
+
+        monkeypatch.setenv("MY_VERSION", "foo")
+        monkeypatch.setenv("MY_TAGS", "b,2,3")
+        monkeypatch.setenv("MY_DESCRIPTION", "1")
+
+        result = await run_sync_in_worker_thread(
+            invoke_and_assert,
+            command=f"deploy ./flows/hello.py:my_flow -n test-name -p {work_pool.name}",
+            expected_output_contains=["bar"],
+        )
+        assert result.exit_code == 0
+        assert "An important name/test" in result.output
+
+        deployment = await prefect_client.read_deployment_by_name(
+            "An important name/test-name"
+        )
+
+        assert deployment.name == "test-name"
+        assert deployment.work_pool_name == "test-work-pool"
+        assert deployment.version == "foo"
+        assert deployment.tags == ["b", ",", "2", ",", "3"]
+        assert deployment.description == "1"
+
     @pytest.mark.usefixtures("project_dir")
     async def test_project_deploy_with_default_parameters(
         self, prefect_client, work_pool

tests/deployment/test_steps.py

@@ -66,6 +66,23 @@ async def test_run_step_resolves_block_document_references_before_running(self):
             "stderr": "",
         }
 
+    async def test_run_step_resolves_environment_variables_before_running(
+        self, monkeypatch
+    ):
+        monkeypatch.setenv("TEST_ENV_VAR", "test_value")
+        output = await run_step(
+            {
+                "prefect.deployments.steps.run_shell_script": {
+                    "script": 'echo "{{ $TEST_ENV_VAR }}"',
+                }
+            }
+        )
+        assert isinstance(output, dict)
+        assert output == {
+            "stdout": "test_value",
+            "stderr": "",
+        }
+
     async def test_run_step_resolves_variables_before_running(self, variables):
         output = await run_step(
             {

tests/utilities/test_templating.py

@@ -81,6 +81,50 @@ def test_finds_block_document_placeholders(self):
         assert placeholder.name == "prefect.blocks.document.name"
         assert placeholder.type is PlaceholderType.BLOCK_DOCUMENT
 
+    def test_finds_env_var_placeholders(self, monkeypatch):
+        monkeypatch.setenv("MY_ENV_VAR", "VALUE")
+        template = "Hello {{$MY_ENV_VAR}}!"
+        placeholders = find_placeholders(template)
+        assert len(placeholders) == 1
+        placeholder = placeholders.pop()
+        assert placeholder.name == "$MY_ENV_VAR"
+        assert placeholder.type is PlaceholderType.ENV_VAR
+
+    def test_apply_values_clears_placeholder_for_missing_env_vars(self):
+        template = "{{ $MISSING_ENV_VAR }}"
+        values = {"ANOTHER_ENV_VAR": "test_value"}
+        result = apply_values(template, values)
+        assert result == ""
+
+    def test_finds_nested_env_var_placeholders(self, monkeypatch):
+        monkeypatch.setenv("GREETING", "VALUE")
+        template = {"greeting": "Hello {{name}}!", "message": {"text": "{{$GREETING}}"}}
+        placeholders = find_placeholders(template)
+        assert len(placeholders) == 2
+        names = set(p.name for p in placeholders)
+        assert names == {"name", "$GREETING"}
+
+        types = set(p.type for p in placeholders)
+        assert types == {PlaceholderType.STANDARD, PlaceholderType.ENV_VAR}
+
+    @pytest.mark.parametrize(
+        "template,expected",
+        [
+            (
+                '{"greeting": "Hello {{name}}!", "message": {"text": "{{$$}}"}}',
+                '{"greeting": "Hello Dan!", "message": {"text": ""}}',
+            ),
+            (
+                '{"greeting": "Hello {{name}}!", "message": {"text": "{{$GREETING}}"}}',
+                '{"greeting": "Hello Dan!", "message": {"text": ""}}',
+            ),
+        ],
+    )
+    def test_invalid_env_var_placeholder(self, template, expected):
+        values = {"name": "Dan"}
+        result = apply_values(template, values)
+        assert result == expected
+
 
 class TestApplyValues:
     def test_apply_values_simple_string_with_one_placeholder(self):
@@ -422,6 +466,7 @@ async def test_resolve_does_not_template_other_placeholder_types(
         template = {
             "key": "{{ another_placeholder }}",
             "key2": "{{ prefect.blocks.arbitraryblock.arbitrary-block }}",
+            "key3": "{{ $another_placeholder }}",
         }
         result = await resolve_variables(template, client=prefect_client)
         assert result == template