feat(tg): write TG_METADATA_DB_USER/PASSWORD into the compose file

dk-installer.py

@@ -2125,6 +2125,24 @@ def find_in_block(contents: str, block: str, key: str) -> typing.Optional[re.Mat
     return re.compile(rf"^([ \t]+){re.escape(key)}:.*$", flags=re.M).search(contents, header.end(), end)
 
 
+def get_testgen_credentials_from_compose(contents: str) -> tuple[typing.Optional[str], typing.Optional[str]]:
+    """Read the bootstrap ``TESTGEN_USERNAME``/``PASSWORD`` out of a compose file's text.
+
+    That pair is the one Postgres account every generated TestGen compose file has always
+    had; ``TG_METADATA_DB_USER/PASSWORD`` is written as a copy of it, never generated fresh.
+    """
+    username = None
+    password = None
+    for line in contents.split("\n"):
+        if line.strip().startswith("TESTGEN_USERNAME:"):
+            username = line.replace("TESTGEN_USERNAME:", "").strip()
+        if line.strip().startswith("TESTGEN_PASSWORD:"):
+            password = line.replace("TESTGEN_PASSWORD:", "").strip()
+        if username and password:
+            break
+    return username, password
+
+
 class UpdateComposeFileStep(Step):
     label = "Updating the Docker compose file"
 
@@ -2135,6 +2153,8 @@ def __init__(self):
         self.update_base_url = False
         self.update_api_port = False
         self.update_stop_grace = False
+        self.update_metadata_creds = False
+        self._missing_metadata_creds_keys = []
         super().__init__()
 
     def pre_execute(self, action, args):
@@ -2207,6 +2227,24 @@ def pre_execute(self, action, args):
             engine_image is not None and find_in_block(contents, "engine", "stop_grace_period") is None
         )
 
+        # TG_METADATA_DB_USER/PASSWORD replaces an implicit app-side fallback to
+        # TESTGEN_USERNAME/PASSWORD, with no fallback of its own. Track each half
+        # independently so a file missing only one of the two still gets repaired,
+        # without re-inserting (and duplicating) the half that's already there.
+        self._missing_metadata_creds_keys = [
+            key for key in ("TG_METADATA_DB_USER", "TG_METADATA_DB_PASSWORD") if key not in contents
+        ]
+        self.update_metadata_creds = bool(self._missing_metadata_creds_keys)
+        if self.update_metadata_creds:
+            username, password = get_testgen_credentials_from_compose(contents)
+            anchor_exists = re.search(r"^([ \t]+)TG_METADATA_DB_HOST:.*$", contents, flags=re.M) is not None
+            if not all([username, password]) or not anchor_exists:
+                CONSOLE.msg(
+                    "Unable to determine TESTGEN_USERNAME/PASSWORD from the existing compose file "
+                    f"[{action.get_compose_file_path(args).absolute()}] to set TG_METADATA_DB_USER/PASSWORD."
+                )
+                raise AbortAction
+
         if not any(
             (
                 self.update_version,
@@ -2215,6 +2253,7 @@ def pre_execute(self, action, args):
                 self.update_base_url,
                 self.update_api_port,
                 self.update_stop_grace,
+                self.update_metadata_creds,
             )
         ):
             CONSOLE.msg("No changes will be applied.")
@@ -2229,6 +2268,7 @@ def execute(self, action, args):
                 self.update_base_url,
                 self.update_api_port,
                 self.update_stop_grace,
+                self.update_metadata_creds,
             )
         ):
             raise SkipStep
@@ -2260,6 +2300,13 @@ def execute(self, action, args):
             var = f"\n{match.group(1)}TG_JWT_HASHING_KEY: {str(base64.b64encode(random.randbytes(32)), 'ascii')}"
             contents = contents[0 : match.end()] + match.group(1) + var + contents[match.end() :]
 
+        if self.update_metadata_creds:
+            username, password = get_testgen_credentials_from_compose(contents)
+            values = {"TG_METADATA_DB_USER": username, "TG_METADATA_DB_PASSWORD": password}
+            match = re.search(r"^([ \t]+)TG_METADATA_DB_HOST:.*$", contents, flags=re.M)
+            var = "".join(f"\n{match.group(1)}{key}: {values[key]}" for key in self._missing_metadata_creds_keys)
+            contents = contents[0 : match.end()] + match.group(1) + var + contents[match.end() :]
+
         if self.update_base_url:
             match = re.search(r"^([ \t]+)TG_METADATA_DB_HOST:.*$", contents, flags=re.M)
             var = f"\n{match.group(1)}TG_UI_BASE_URL: {self._base_url}"
@@ -2292,7 +2339,7 @@ def __init__(self):
     def pre_execute(self, action, args):
         super().pre_execute(action, args)
         if action.ctx.get("using_existing"):
-            self.username, self.password = self.get_credentials_from_compose_file(
+            self.username, self.password = get_testgen_credentials_from_compose(
                 action.get_compose_file_path(args).read_text()
             )
         else:
@@ -2321,18 +2368,6 @@ def on_action_success(self, action, args):
 
         CONSOLE.msg(f"(Credentials also written to {simplify_path(cred_file_path)})")
 
-    def get_credentials_from_compose_file(self, file_contents):
-        username = None
-        password = None
-        for line in file_contents.split("\n"):
-            if line.strip().startswith("TESTGEN_USERNAME:"):
-                username = line.replace("TESTGEN_USERNAME:", "").strip()
-            if line.strip().startswith("TESTGEN_PASSWORD:"):
-                password = line.replace("TESTGEN_PASSWORD:", "").strip()
-            if username and password:
-                break
-        return username, password
-
     def get_compose_file_contents(self, action, args):
         action.analytics.additional_properties["used_custom_cert"] = bool(args.ssl_cert_file and args.ssl_key_file)
         action.analytics.additional_properties["used_custom_image"] = args.image != TESTGEN_DEFAULT_IMAGE
@@ -2368,6 +2403,8 @@ def get_compose_file_contents(self, action, args):
               TG_DECRYPT_PASSWORD: {generate_password()}
               TG_JWT_HASHING_KEY: {str(base64.b64encode(random.randbytes(32)), "ascii")}
               TG_METADATA_DB_HOST: postgres
+              TG_METADATA_DB_USER: {self.username}
+              TG_METADATA_DB_PASSWORD: {self.password}
               TG_TARGET_DB_TRUST_SERVER_CERTIFICATE: yes
               TG_EXPORT_TO_OBSERVABILITY_VERIFY_SSL: no
               TG_INSTANCE_ID: {action.analytics.get_instance_id()}

tests/test_tg_install.py

@@ -1,4 +1,5 @@
 import json
+import re
 from functools import partial
 from pathlib import Path
 from unittest.mock import call, patch
@@ -136,6 +137,24 @@ def test_tg_compose_sets_engine_stop_grace_period(tg_install_action, start_cmd_m
     assert lines[grace_idx] == f"{indent}stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s"
 
 
+@pytest.mark.integration
+def test_tg_compose_metadata_db_creds_match_bootstrap_user(
+    tg_install_action, start_cmd_mock, stdout_mock, compose_path
+):
+    """TG_METADATA_DB_USER/PASSWORD replaces an implicit app-side fallback to
+    TESTGEN_USERNAME/PASSWORD — no new Postgres role is created, so both pairs must carry
+    the same bootstrap credentials."""
+    tg_install_action.execute()
+    contents = compose_path.read_text()
+
+    username_match = re.search(r"TESTGEN_USERNAME: (\S+)", contents)
+    password_match = re.search(r"TESTGEN_PASSWORD: (\S+)", contents)
+    username, password = username_match.group(1), password_match.group(1)
+
+    assert f"TG_METADATA_DB_USER: {username}" in contents
+    assert f"TG_METADATA_DB_PASSWORD: {password}" in contents
+
+
 @pytest.mark.integration
 def test_tg_compose_base_url_ssl(tg_install_action, start_cmd_mock, stdout_mock, args_mock, compose_path):
     args_mock.ssl_cert_file = "/path/to/cert.crt"

tests/test_tg_upgrade.py

@@ -12,6 +12,7 @@
     TESTGEN_STOP_GRACE_PERIOD,
     TestgenUpgradeAction,
     find_in_block,
+    get_testgen_credentials_from_compose,
     InstallMarker,
 )
 
@@ -145,7 +146,11 @@ def test_tg_upgrade_abort(
     args_mock.skip_verify = False
     set_version_check_mock(version_check_mock, "1.0.0")
     initial_compose_content = get_compose_content(
-        "TG_INSTANCE_ID: test-instance-id", "TG_UI_BASE_URL: http://localhost:8501", stop_grace=True
+        "TG_INSTANCE_ID: test-instance-id",
+        "TG_UI_BASE_URL: http://localhost:8501",
+        "TG_METADATA_DB_USER: admin",
+        "TG_METADATA_DB_PASSWORD: WOzviKBQJS50",
+        stop_grace=True,
     )
     compose_path.write_text(initial_compose_content)
 
@@ -249,6 +254,117 @@ def test_tg_upgrade_preserves_existing_base_url(
     assert compose_content.count("TG_UI_BASE_URL") == 1
 
 
+@pytest.mark.integration
+def test_tg_upgrade_adds_metadata_db_creds(
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    tg_upgrade_stdout_side_effect,
+    args_mock,
+    version_check_mock,
+):
+    """Existing installs never had TG_METADATA_DB_USER/PASSWORD — the upgrade backfills it
+    from the bootstrap TESTGEN_USERNAME/PASSWORD already in the file, since no new Postgres
+    role is created."""
+    set_version_check_mock(version_check_mock, "1.0.0")
+    compose_path.write_text(get_compose_content("TG_INSTANCE_ID: test-instance-id"))
+
+    tg_upgrade_action.execute(args_mock)
+
+    compose_content = compose_path.read_text()
+    assert "TG_METADATA_DB_USER: admin" in compose_content
+    assert "TG_METADATA_DB_PASSWORD: WOzviKBQJS50" in compose_content
+
+
+@pytest.mark.integration
+def test_tg_upgrade_backfills_only_the_missing_metadata_db_half(
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    tg_upgrade_stdout_side_effect,
+    args_mock,
+    version_check_mock,
+):
+    """A file that already has TG_METADATA_DB_USER (e.g. from a prior backfill) but is
+    missing TG_METADATA_DB_PASSWORD must get only the password added — re-inserting the
+    user too would duplicate the key."""
+    set_version_check_mock(version_check_mock, "1.0.0")
+    compose_path.write_text(get_compose_content("TG_INSTANCE_ID: test-instance-id", "TG_METADATA_DB_USER: admin"))
+
+    tg_upgrade_action.execute(args_mock)
+
+    compose_content = compose_path.read_text()
+    assert "TG_METADATA_DB_PASSWORD: WOzviKBQJS50" in compose_content
+    assert compose_content.count("TG_METADATA_DB_USER") == 1
+
+
+@pytest.mark.integration
+def test_tg_upgrade_aborts_when_bootstrap_creds_unavailable(
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    args_mock,
+    console_msg_mock,
+):
+    """TESTGEN_USERNAME/PASSWORD missing or hand-edited out of the file means there's
+    nothing to backfill TG_METADATA_DB_USER/PASSWORD from — abort rather than writing
+    literal 'None' values into a var the app now requires with no fallback."""
+    args_mock.skip_verify = True
+    initial_compose_content = textwrap.dedent("""
+        name: testgen
+
+        x-common-variables: &common-variables
+          TG_DECRYPT_SALT: zyIJQsuBImx5
+          TG_DECRYPT_PASSWORD: cAEGUVRwxvVg
+          TG_JWT_HASHING_KEY: VGVzdEdlbgo=
+          TG_METADATA_DB_HOST: postgres
+          TG_TARGET_DB_TRUST_SERVER_CERTIFICATE: yes
+          TG_EXPORT_TO_OBSERVABILITY_VERIFY_SSL: no
+          TG_INSTANCE_ID: test-instance-id
+          TG_UI_BASE_URL: http://localhost:8501
+
+        services:
+          engine:
+            image: datakitchen/dataops-testgen:v2.14.5
+            stop_grace_period: 90s
+    """)
+    compose_path.write_text(initial_compose_content)
+
+    with pytest.raises(AbortAction):
+        tg_upgrade_action.execute(args_mock)
+
+    console_msg_mock.assert_any_msg_contains("Unable to determine TESTGEN_USERNAME/PASSWORD")
+    assert compose_path.read_text() == initial_compose_content
+    start_cmd_mock.assert_not_called()
+
+
+@pytest.mark.integration
+def test_tg_upgrade_preserves_existing_metadata_db_creds(
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    tg_upgrade_stdout_side_effect,
+    args_mock,
+    version_check_mock,
+):
+    args_mock.skip_verify = True
+    set_version_check_mock(version_check_mock, "1.1.0")
+    compose_path.write_text(
+        get_compose_content(
+            "TG_INSTANCE_ID: test-instance-id",
+            "TG_UI_BASE_URL: https://custom.example.com",
+            "TG_METADATA_DB_USER: custom-user",
+            "TG_METADATA_DB_PASSWORD: custom-pass",
+        )
+    )
+
+    tg_upgrade_action.execute(args_mock)
+
+    compose_content = compose_path.read_text()
+    assert "TG_METADATA_DB_USER: custom-user" in compose_content
+    assert compose_content.count("TG_METADATA_DB_USER") == 1
+
+
 @pytest.mark.integration
 def test_tg_upgrade_adds_stop_grace_period(
     tg_upgrade_action,
@@ -402,3 +518,17 @@ def test_find_in_block_offsets_are_absolute():
     match = find_in_block(COMPOSE_TWO_SERVICES, "postgres", "image")
     assert COMPOSE_TWO_SERVICES[match.start() : match.end()] == "    image: postgres:14.1-alpine"
     assert match.group(1) == "    "
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "contents, expected",
+    (
+        ("TESTGEN_USERNAME: admin\nTESTGEN_PASSWORD: secret\n", ("admin", "secret")),
+        ("  TESTGEN_USERNAME: admin\n  TESTGEN_PASSWORD: secret\n", ("admin", "secret")),
+        ("TESTGEN_PASSWORD: secret\n", (None, "secret")),
+        ("", (None, None)),
+    ),
+)
+def test_get_testgen_credentials_from_compose(contents, expected):
+    assert get_testgen_credentials_from_compose(contents) == expected