Skip to content

Commit

Permalink
Update Settings.copy_with_update to ignore sources when restoring d…
Browse files Browse the repository at this point in the history
…efaults
  • Loading branch information
desertaxle committed Oct 31, 2024
1 parent 5cfdfd3 commit ef64af8
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 2 deletions.
28 changes: 26 additions & 2 deletions src/prefect/settings/models/root.py
Original file line number Diff line number Diff line change
Expand Up @@ -290,9 +290,32 @@ def copy_with_update(
Returns:
A new Settings object.
"""
# To restore defaults, we need to resolve the setting path and then
# set the default value on the new settings object. When restoring
# defaults, all settings sources will be ignored.
restore_defaults_obj = {}
for r in restore_defaults or []:
set_in_dict(restore_defaults_obj, r.accessor, True)
path = r.accessor.split(".")
model = self
for key in path[:-1]:
model = model.model_fields[key].annotation
assert model is not None, f"Invalid setting path: {r.accessor}"

model_field = model.model_fields[path[-1]]
assert model is not None, f"Invalid setting path: {r.accessor}"
if hasattr(model_field, "default"):
default = model_field.default
elif (
hasattr(model_field, "default_factory") and model_field.default_factory
):
default = model_field.default_factory()
else:
raise ValueError(f"No default value for setting: {r.accessor}")
set_in_dict(
restore_defaults_obj,
r.accessor,
default,
)
updates = updates or {}
set_defaults = set_defaults or {}

Expand All @@ -307,7 +330,8 @@ def copy_with_update(
new_settings = self.__class__.model_validate(
deep_merge_dicts(
set_defaults_obj,
self.model_dump(exclude_unset=True, exclude=restore_defaults_obj),
self.model_dump(exclude_unset=True),
restore_defaults_obj,
updates_obj,
)
)
Expand Down
9 changes: 9 additions & 0 deletions tests/test_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -532,6 +532,15 @@ def test_settings_copy_with_update(self):
), "Changed, existing value was default"
assert new_settings.client.retry_extra_codes == {400, 500}

def test_settings_copy_with_update_restore_defaults(self, monkeypatch):
monkeypatch.setenv("PREFECT_TESTING_TEST_SETTING", "Not the default")
settings = Settings()
assert settings.testing.test_setting == "Not the default"
new_settings = settings.copy_with_update(
restore_defaults={PREFECT_TEST_SETTING},
)
assert new_settings.testing.test_setting == "FOO"

def test_settings_loads_environment_variables_at_instantiation(self, monkeypatch):
assert PREFECT_TEST_MODE.value() is True

Expand Down

0 comments on commit ef64af8

Please sign in to comment.