diff --git a/awx/lib/awx_display_callback/module.py b/awx/lib/awx_display_callback/module.py index 6800560cfc..368063d0d1 100644 --- a/awx/lib/awx_display_callback/module.py +++ b/awx/lib/awx_display_callback/module.py @@ -18,7 +18,11 @@ from __future__ import (absolute_import, division, print_function) # Python +import codecs import contextlib +import json +import os +import stat import sys import uuid from copy import copy @@ -292,10 +296,22 @@ class BaseCallbackModule(CallbackBase): failures=stats.failures, ok=stats.ok, processed=stats.processed, - skipped=stats.skipped, - artifact_data=stats.custom.get('_run', {}) if hasattr(stats, 'custom') else {} + skipped=stats.skipped ) + # write custom set_stat artifact data to the local disk so that it can + # be persisted by awx after the process exits + custom_artifact_data = stats.custom.get('_run', {}) if hasattr(stats, 'custom') else {} + if custom_artifact_data: + # create the directory for custom stats artifacts to live in (if it doesn't exist) + custom_artifacts_dir = os.path.join(os.getenv('AWX_PRIVATE_DATA_DIR'), 'artifacts') + os.makedirs(custom_artifacts_dir, mode=stat.S_IXUSR + stat.S_IWUSR + stat.S_IRUSR) + + custom_artifacts_path = os.path.join(custom_artifacts_dir, 'custom') + with codecs.open(custom_artifacts_path, 'w', encoding='utf-8') as f: + os.chmod(custom_artifacts_path, stat.S_IRUSR | stat.S_IWUSR) + json.dump(custom_artifact_data, f) + with self.capture_event_data('playbook_on_stats', **event_data): super(BaseCallbackModule, self).v2_playbook_on_stats(stats) diff --git a/awx/lib/tests/test_display_callback.py b/awx/lib/tests/test_display_callback.py index fefd5d4188..e87f3ec306 100644 --- a/awx/lib/tests/test_display_callback.py +++ b/awx/lib/tests/test_display_callback.py @@ -2,7 +2,9 @@ from collections import OrderedDict import json import mock import os +import shutil import sys +import tempfile import pytest @@ -254,3 +256,26 @@ def test_callback_plugin_strips_task_environ_variables(executor, cache, playbook assert len(cache) for event in cache.values(): assert os.environ['PATH'] not in json.dumps(event) + + +@pytest.mark.parametrize('playbook', [ +{'custom_set_stat.yml': ''' +- name: custom set_stat calls should persist to the local disk so awx can save them + connection: local + hosts: all + tasks: + - set_stats: + data: + foo: "bar" +'''}, # noqa +]) +def test_callback_plugin_saves_custom_stats(executor, cache, playbook): + try: + private_data_dir = tempfile.mkdtemp() + with mock.patch.dict(os.environ, {'AWX_PRIVATE_DATA_DIR': private_data_dir}): + executor.run() + artifacts_path = os.path.join(private_data_dir, 'artifacts', 'custom') + with open(artifacts_path, 'r') as f: + assert json.load(f) == {'foo': 'bar'} + finally: + shutil.rmtree(os.path.join(private_data_dir)) diff --git a/awx/main/tasks.py b/awx/main/tasks.py index 3345dada44..9b04d368b8 100644 --- a/awx/main/tasks.py +++ b/awx/main/tasks.py @@ -659,6 +659,7 @@ class BaseTask(LogErrorsTask): # Derived class should call add_ansible_venv() or add_awx_venv() if self.should_use_proot(instance, **kwargs): env['PROOT_TMP_DIR'] = settings.AWX_PROOT_BASE_PATH + env['AWX_PRIVATE_DATA_DIR'] = kwargs['private_data_dir'] return env def build_safe_env(self, env, **kwargs): @@ -1307,6 +1308,21 @@ class RunJob(BaseTask): kwargs['private_data_dir'], kwargs['fact_modification_times'] ) + + # persist artifacts set via `set_stat` (if any) + custom_stats_path = os.path.join(kwargs['private_data_dir'], 'artifacts', 'custom') + if os.path.exists(custom_stats_path): + with open(custom_stats_path, 'r') as f: + custom_stat_data = None + try: + custom_stat_data = json.load(f) + except ValueError: + logger.warning('Could not parse custom `set_fact` data for job {}'.format(job.id)) + + if custom_stat_data: + job.artifacts = custom_stat_data + job.save(update_fields=['artifacts']) + try: inventory = job.inventory except Inventory.DoesNotExist: