From fc11d1601ea2bf266019d0d63d20521df2a731e1 Mon Sep 17 00:00:00 2001
From: Peter Kraft <petereliaskraft@gmail.com>
Date: Wed, 31 Jul 2024 12:15:46 -0700
Subject: [PATCH] Envsubst (#20)

---
 dbos_transact/dbos_config.py     | 18 +++++++++++++++++-
 templates/hello/dbos-config.yaml |  2 +-
 tests/test_config.py             |  9 +++++++--
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/dbos_transact/dbos_config.py b/dbos_transact/dbos_config.py
index 37c87641..b659dc05 100644
--- a/dbos_transact/dbos_config.py
+++ b/dbos_transact/dbos_config.py
@@ -1,4 +1,6 @@
 import json
+import os
+import re
 from importlib import resources
 from typing import Any, Dict, List, Optional, TypedDict
 
@@ -47,10 +49,24 @@ class ConfigFile(TypedDict):
     env: Dict[str, str]
 
 
+def substitute_env_vars(content: str) -> str:
+    regex = r"\$\{([^}]+)\}"  # Regex to match ${VAR_NAME} style placeholders
+
+    def replace_func(match: re.Match[str]) -> str:
+        var_name = match.group(1)
+        return os.environ.get(
+            var_name, '""'
+        )  # If the env variable is not set, return an empty string
+
+    return re.sub(regex, replace_func, content)
+
+
 def load_config(configFilePath: str = "dbos-config.yaml") -> ConfigFile:
     # Load the YAML file
     with open(configFilePath, "r") as file:
-        data = yaml.safe_load(file)
+        content = file.read()
+        substituted_content = substitute_env_vars(content)
+        data = yaml.safe_load(substituted_content)
 
     # Load the JSON schema relative to the package root
     with resources.open_text("dbos_transact", "dbos-config.schema.json") as schema_file:
diff --git a/templates/hello/dbos-config.yaml b/templates/hello/dbos-config.yaml
index ccd6b63e..ed19a077 100644
--- a/templates/hello/dbos-config.yaml
+++ b/templates/hello/dbos-config.yaml
@@ -10,7 +10,7 @@ database:
   hostname: localhost
   port: 5432
   username: postgres
-  password: dbos
+  password: ${PGPASSWORD}
   app_db_name: hello
 telemetry:
   logs:
diff --git a/tests/test_config.py b/tests/test_config.py
index ad3707df..1e4bd850 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,5 +1,6 @@
 # type: ignore
 
+import os
 from unittest.mock import mock_open
 
 import pytest
@@ -27,10 +28,13 @@ def test_valid_config(mocker):
           hostname: 'some host'
           port: 1234
           username: 'some user'
-          password: abc123
+          password: ${PGPASSWORD}
           app_db_name: 'some db'
           connectionTimeoutMillis: 3000
+        env:
+            foo: ${BARBAR}
     """
+    os.environ["BARBAR"] = "FOOFOO"
     mocker.patch(
         "builtins.open", side_effect=generate_mock_open(mock_filename, mock_config)
     )
@@ -39,9 +43,10 @@ def test_valid_config(mocker):
     assert configFile["database"]["hostname"] == "some host"
     assert configFile["database"]["port"] == 1234
     assert configFile["database"]["username"] == "some user"
-    assert configFile["database"]["password"] == "abc123"
+    assert configFile["database"]["password"] == os.environ["PGPASSWORD"]
     assert configFile["database"]["app_db_name"] == "some db"
     assert configFile["database"]["connectionTimeoutMillis"] == 3000
+    assert configFile["env"]["foo"] == "FOOFOO"
 
 
 def test_config_missing_params(mocker):