diff --git a/dk-installer.py b/dk-installer.py index 53c5782..c23fac9 100755 --- a/dk-installer.py +++ b/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()} diff --git a/tests/test_tg_install.py b/tests/test_tg_install.py index d4d9b87..d4566f8 100644 --- a/tests/test_tg_install.py +++ b/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" diff --git a/tests/test_tg_upgrade.py b/tests/test_tg_upgrade.py index 21d57d6..9d22026 100644 --- a/tests/test_tg_upgrade.py +++ b/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