From 198f2e31a44fdace6c04ce1472dfb144ffaf39ce Mon Sep 17 00:00:00 2001 From: Ryan Petrello Date: Wed, 19 Jul 2017 14:23:20 -0400 Subject: [PATCH] consolidate tasks.py cleanup for temporary files into a single place see: #6199 --- awx/main/tasks.py | 61 +++++++++++++------------------ awx/main/tests/unit/test_tasks.py | 11 ++++++ 2 files changed, 37 insertions(+), 35 deletions(-) diff --git a/awx/main/tasks.py b/awx/main/tasks.py index 4624435e36..c8ac70f031 100644 --- a/awx/main/tasks.py +++ b/awx/main/tasks.py @@ -6,6 +6,7 @@ import codecs from collections import OrderedDict import ConfigParser import cStringIO +import functools import imp import json import logging @@ -382,10 +383,29 @@ def delete_inventory(inventory_id): return +def with_path_cleanup(f): + @functools.wraps(f) + def _wrapped(self, *args, **kwargs): + try: + return f(self, *args, **kwargs) + finally: + for p in self.cleanup_paths: + try: + if os.path.isdir(p): + shutil.rmtree(p, ignore_errors=True) + elif os.path.exists(p): + os.remove(p) + except OSError: + logger.exception("Failed to remove tmp file: {}".format(p)) + self.cleanup_paths = [] + return _wrapped + + class BaseTask(Task): name = None model = None abstract = True + cleanup_paths = [] def update_model(self, pk, _attempt=0, **updates): """Reload the model instance from the database and update the @@ -454,6 +474,7 @@ class BaseTask(Task): ''' path = tempfile.mkdtemp(prefix='ansible_tower_%s_' % instance.pk, dir=settings.AWX_PROOT_BASE_PATH) os.chmod(path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) + self.cleanup_paths.append(path) return path def build_private_data_files(self, instance, **kwargs): @@ -666,6 +687,7 @@ class BaseTask(Task): Hook for any steps to run after job/task is marked as complete. ''' + @with_path_cleanup def run(self, pk, isolated_host=None, **kwargs): ''' Run the job/task and capture its output. @@ -730,6 +752,7 @@ class BaseTask(Task): if not check_proot_installed(): raise RuntimeError('bubblewrap is not installed') kwargs['proot_temp_dir'] = build_proot_temp_dir() + self.cleanup_paths.append(kwargs['proot_temp_dir']) args = wrap_args_with_proot(args, cwd, **kwargs) safe_args = wrap_args_with_proot(safe_args, cwd, **kwargs) # If there is an SSH key path defined, wrap args with ssh-agent. @@ -779,16 +802,6 @@ class BaseTask(Task): if settings.DEBUG: logger.exception('exception occurred while running task') finally: - if kwargs.get('private_data_dir', ''): - try: - shutil.rmtree(kwargs['private_data_dir'], True) - except OSError: - pass - if kwargs.get('proot_temp_dir', ''): - try: - shutil.rmtree(kwargs['proot_temp_dir'], True) - except OSError: - pass try: stdout_handle.flush() stdout_handle.close() @@ -1215,6 +1228,7 @@ class RunProjectUpdate(BaseTask): } ''' handle, self.revision_path = tempfile.mkstemp(dir=settings.AWX_PROOT_BASE_PATH) + self.cleanup_paths.append(self.revision_path) private_data = {'credentials': {}} if project_update.credential: credential = project_update.credential @@ -1498,13 +1512,6 @@ class RunProjectUpdate(BaseTask): if status == 'successful' and instance.launch_type != 'sync': self._update_dependent_inventories(instance, dependent_inventory_sources) - def final_run_hook(self, instance, status, **kwargs): - super(RunProjectUpdate, self).final_run_hook(instance, status, **kwargs) - try: - os.remove(self.revision_path) - except Exception, e: - logger.error("Failed removing revision tmp file: {}".format(e)) - class RunInventoryUpdate(BaseTask): @@ -1841,7 +1848,7 @@ class RunInventoryUpdate(BaseTask): elif src == 'scm': args.append(inventory_update.get_actual_source_path()) elif src == 'custom': - runpath = tempfile.mkdtemp(prefix='ansible_tower_launch_', dir=settings.AWX_PROOT_BASE_PATH) + runpath = tempfile.mkdtemp(prefix='ansible_tower_inventory_', dir=settings.AWX_PROOT_BASE_PATH) handle, path = tempfile.mkstemp(dir=runpath) f = os.fdopen(handle, 'w') if inventory_update.source_script is None: @@ -1851,7 +1858,7 @@ class RunInventoryUpdate(BaseTask): os.chmod(path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) args.append(path) args.append("--custom") - self.custom_dir_path.append(runpath) + self.cleanup_paths.append(runpath) args.append('-v%d' % inventory_update.verbosity) if settings.DEBUG: args.append('--traceback') @@ -1873,8 +1880,6 @@ class RunInventoryUpdate(BaseTask): return getattr(settings, 'INVENTORY_UPDATE_IDLE_TIMEOUT', None) def pre_run_hook(self, inventory_update, **kwargs): - self.custom_dir_path = [] - source_project = None if inventory_update.inventory_source: source_project = inventory_update.inventory_source.source_project @@ -1906,14 +1911,6 @@ class RunInventoryUpdate(BaseTask): ('project_update', local_project_sync.name, local_project_sync.id))) raise - def final_run_hook(self, instance, status, **kwargs): - if self.custom_dir_path: - for p in self.custom_dir_path: - try: - shutil.rmtree(p, True) - except OSError: - pass - class RunAdHocCommand(BaseTask): ''' @@ -2106,12 +2103,6 @@ class RunAdHocCommand(BaseTask): ''' return getattr(settings, 'AWX_PROOT_ENABLED', False) - def final_run_hook(self, ad_hoc_command, status, **kwargs): - ''' - Hook for actions to run after ad hoc command is marked as completed. - ''' - super(RunAdHocCommand, self).final_run_hook(ad_hoc_command, status, **kwargs) - class RunSystemJob(BaseTask): diff --git a/awx/main/tests/unit/test_tasks.py b/awx/main/tests/unit/test_tasks.py index dad6429fa0..162904629c 100644 --- a/awx/main/tests/unit/test_tasks.py +++ b/awx/main/tests/unit/test_tasks.py @@ -262,6 +262,17 @@ class TestGenericRun(TestJobExecution): ]: assert c in self.task.update_model.call_args_list + def test_artifact_cleanup(self): + path = tempfile.NamedTemporaryFile(delete=False).name + try: + self.task.cleanup_paths.append(path) + assert os.path.exists(path) + self.task.run(self.pk) + assert not os.path.exists(path) + finally: + if os.path.exists(path): + os.remove(path) + def test_uses_bubblewrap(self): self.task.run(self.pk)