Skip to content

Commit

Permalink
[AIRFLOW-6060] Improve conf_vars context manager (apache#6658)
Browse files Browse the repository at this point in the history
This commit adds try / finally clause to conf_vars context
manager to assure that initail values are reseted in case
of an exception in yield.
  • Loading branch information
turbaszek authored and potiuk committed Nov 25, 2019
1 parent ae96f27 commit 27ce7bd
Showing 1 changed file with 11 additions and 15 deletions.
26 changes: 11 additions & 15 deletions tests/test_utils/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@
def conf_vars(overrides):
original = {}
original_env_vars = {}
reconfigure_vars = False
for (section, key), value in overrides.items():

env = conf._env_var_name(section, key)
Expand All @@ -43,18 +42,15 @@ def conf_vars(overrides):
conf.set(section, key, value)
else:
conf.remove_option(section, key)

if section == 'core' and key.lower().endswith('_folder'):
reconfigure_vars = True
if reconfigure_vars:
settings.configure_vars()
yield
for (section, key), value in original.items():
if value is not None:
conf.set(section, key, value)
else:
conf.remove_option(section, key)
for env, value in original_env_vars.items():
os.environ[env] = value
if reconfigure_vars:
settings.configure_vars()
try:
yield
finally:
for (section, key), value in original.items():
if value is not None:
conf.set(section, key, value)
else:
conf.remove_option(section, key)
for env, value in original_env_vars.items():
os.environ[env] = value
settings.configure_vars()

0 comments on commit 27ce7bd

Please sign in to comment.