Repository: bentoml/BentoML
Pull request #4254 (2023-11-02 02:34:44)

feat: support storing secrets with env vars in config
Signed-off-by: Frost Ming <me@frostming.com>

## What does this PR address?

<!--
Thanks for sending a pull request!

Congrats for making it this far! Here's a 🍱 for you. There are still a few steps ahead.

Please make sure to read the contribution guidelines, then fill out the blanks below before requesting a code review.

Name your Pull Request with one of the following prefixes, e.g. "feat: add support for PyTorch", to indicate the type of changes proposed. This is based on the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary).
  - feat: (new feature for the user, not a new feature for build script)
  - fix: (bug fix for the user, not a fix to a build script)
  - docs: (changes to the documentation)
  - style: (formatting, missing semicolons, etc; no production code change)
  - refactor: (refactoring production code, eg. renaming a variable)
  - perf: (code changes that improve performance)
  - test: (adding missing tests, refactoring tests; no production code change)
  - chore: (updating grunt tasks etc; no production code change)
  - build: (changes that affect the build system or external dependencies)
  - ci: (changes to configuration files and scripts)
  - revert: (reverts a previous commit)

Describe your changes in detail. Attach screenshots here if appropriate.

Once you're done with this, someone from BentoML team or community member will help review your PR (see "Who can help review?" section for potential reviewers.). If no one has reviewed your PR after a week have passed, don't hesitate to post a new comment and ping @-the same person. Notifications sometimes get lost 🥲.
-->

<!-- Remove if not applicable -->

Fixes #(issue)

## Before submitting:

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
<!--- If you plan to update documentation or tests in follow-up, please note -->

- [ ] Does the Pull Request follow [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary) naming? Here are [GitHub's
      guide](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request) on how to create a pull request.
- [ ] Does the code follow BentoML's code style, `pre-commit run -a` script has passed ([instructions](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#style-check-auto-formatting-type-checking))?
- [ ] Did you read through [contribution guidelines](https://github.com/bentoml/BentoML/blob/main/CONTRIBUTING.md#ways-to-contribute) and follow [development guidelines](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#start-developing)?
- [ ] Did your changes require updates to the documentation? Have you updated
      those accordingly? Here are [documentation guidelines](https://github.com/bentoml/BentoML/tree/main/docs) and [tips on writting docs](https://github.com/bentoml/BentoML/tree/main/docs#writing-documentation).
- [ ] Did you write tests to cover your changes?



Patch:
diff --git a/docs/source/guides/configuration.rst b/docs/source/guides/configuration.rst
index af19a57d..501c1026 100644
--- a/docs/source/guides/configuration.rst
+++ b/docs/source/guides/configuration.rst
@@ -55,6 +55,23 @@ below:
    In the example above, the number of API workers count is overridden to 4.
    Remaining properties will take their defaults values.
 
+Variables in the form of ``${ENV_VAR}`` will be expanded at runtime to the value of the corresponding environment variable, but please note that this only supports string types.
+For example:
+
+.. code-block:: yaml
+   :caption: `~/bentoml_configuration.yaml`
+
+   ssl:
+     keyfile_password: ${MY_SSL_KEYFILE_PASSWORD}
+
+In addition, you can provide default values that will take effect when the environment variable is not set in the following form：
+
+.. code-block:: yaml
+   :caption: `~/bentoml_configuration.yaml`
+
+   ssl:
+     keyfile_password: ${MY_SSL_KEYFILE_PASSWORD:-default_value}
+
 .. seealso::
 
    :ref:`guides/configuration:Configuration fields`
diff --git a/src/bentoml/_internal/configuration/containers.py b/src/bentoml/_internal/configuration/containers.py
index 068a783c..99139a50 100644
--- a/src/bentoml/_internal/configuration/containers.py
+++ b/src/bentoml/_internal/configuration/containers.py
@@ -25,6 +25,7 @@ from ..utils import split_with_quotes
 from ..utils import validate_or_create_dir
 from ..utils.unflatten import unflatten
 from . import expand_env_var
+from .helpers import expand_env_var_in_values
 from .helpers import flatten_dict
 from .helpers import get_default_config
 from .helpers import import_configuration_spec
@@ -166,6 +167,7 @@ class BentoMLConfiguration:
                     deepcopy(global_runner_cfg),
                     runner_cfg,
                 )
+        expand_env_var_in_values(self.config)
 
     def to_dict(self) -> providers.ConfigDictType:
         return t.cast(providers.ConfigDictType, self.config)
diff --git a/src/bentoml/_internal/configuration/helpers.py b/src/bentoml/_internal/configuration/helpers.py
index 4603eb2a..c100c02e 100644
--- a/src/bentoml/_internal/configuration/helpers.py
+++ b/src/bentoml/_internal/configuration/helpers.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import ipaddress
 import logging
 import os
+import re
 import typing as t
 from functools import singledispatch
 from typing import TYPE_CHECKING
@@ -165,3 +166,28 @@ def is_valid_ip_address(addr: str) -> bool:
         return True
     except ValueError:
         return False
+
+
+_ENV_VAR_PATTERN = r"\$\{(?P<name>.+?)(?:\:-(?P<default>.+?))?\}"
+
+
+def expand_env_var(value: t.Any) -> str:
+    if not isinstance(value, str) or "${" not in value:
+        return value
+
+    def _replace_env(match: t.Match[str]) -> str:
+        name = match.group("name")
+        default = match.group("default")
+        return os.getenv(name, default or "")
+
+    return re.sub(_ENV_VAR_PATTERN, _replace_env, value)
+
+
+def expand_env_var_in_values(d: t.MutableMapping[str, t.Any]) -> None:
+    for k, v in d.items():
+        if isinstance(v, t.MutableMapping):
+            expand_env_var_in_values(v)
+        elif isinstance(v, str):
+            d[k] = expand_env_var(v)
+        elif isinstance(v, t.Sequence):
+            d[k] = [expand_env_var(i) for i in v]


Test patch:
diff --git a/tests/unit/_internal/configuration/test_containers.py b/tests/unit/_internal/configuration/test_containers.py
index 2ac60c88..4965584d 100644
--- a/tests/unit/_internal/configuration/test_containers.py
+++ b/tests/unit/_internal/configuration/test_containers.py
@@ -161,3 +161,17 @@ runners:
     assert runner_cfg["test_runner_1"]["traffic"]["timeout"] == 100
     assert runner_cfg["test_runner_2"]["traffic"]["timeout"] == 50
     assert runner_cfg["test_runner_3"]["traffic"]["timeout"] == 60
+
+
+def test_expand_env_vars_in_values(container_from_file, monkeypatch):
+    CONFIG = """\
+api_server:
+    ssl:
+        keyfile_password: ${KEYFILE_PASSWORD:-123456}
+"""
+    config = container_from_file(CONFIG)
+    assert config["api_server"]["ssl"]["keyfile_password"] == "123456"
+
+    monkeypatch.setenv("KEYFILE_PASSWORD", "654321")
+    config = container_from_file(CONFIG)
+    assert config["api_server"]["ssl"]["keyfile_password"] == "654321"
diff --git a/tests/unit/_internal/configuration/test_helpers.py b/tests/unit/_internal/configuration/test_helpers.py
index e9b908ae..ff27d21f 100644
--- a/tests/unit/_internal/configuration/test_helpers.py
+++ b/tests/unit/_internal/configuration/test_helpers.py
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 
 import pytest
 
+from bentoml._internal.configuration.helpers import expand_env_var_in_values
 from bentoml._internal.configuration.helpers import flatten_dict
 from bentoml._internal.configuration.helpers import is_valid_ip_address
 from bentoml._internal.configuration.helpers import load_config_file
@@ -134,3 +135,48 @@ def test_invalid_ip_address():
     assert not is_valid_ip_address("256.255.255.255")
     assert not is_valid_ip_address("256.256.256.256")
     assert not is_valid_ip_address("")
+
+
+def test_expand_env_in_values(monkeypatch: pytest.MonkeyPatch):
+    data = {
+        "foo": "${FOO}",
+        "foo-list": ["${FOO_ITEM1:-default1}", "${FOO_ITEM2:-default2}"],
+        "bar": {
+            "name": "${BAR_NAME:-bardefault}",
+            "bar-list": ["${BAR_ITEM1:-default1}", "${BAR_ITEM2:-default2}"],
+        },
+    }
+    expand_env_var_in_values(data)
+    assert data == {
+        "foo": "",
+        "foo-list": ["default1", "default2"],
+        "bar": {
+            "name": "bardefault",
+            "bar-list": ["default1", "default2"],
+        },
+    }
+    data = {
+        "foo": "${FOO}",
+        "foo-list": ["${FOO_ITEM1:-default1}", "${FOO_ITEM2:-default2}"],
+        "bar": {
+            "name": "${BAR_NAME:-bardefault}",
+            "bar-list": ["${BAR_ITEM1:-default1}", "${BAR_ITEM2:-default2}"],
+        },
+    }
+    with monkeypatch.context() as m:
+        m.setenv("FOO", "foo")
+        m.setenv("FOO_ITEM1", "foo-item1")
+        m.setenv("FOO_ITEM2", "foo-item2")
+        m.setenv("BAR_NAME", "bar-name")
+        m.setenv("BAR_ITEM1", "bar-item1")
+        m.setenv("BAR_ITEM2", "bar-item2")
+
+        expand_env_var_in_values(data)
+        assert data == {
+            "foo": "foo",
+            "foo-list": ["foo-item1", "foo-item2"],
+            "bar": {
+                "name": "bar-name",
+                "bar-list": ["bar-item1", "bar-item2"],
+            },
+        }