diff --git a/ansible_mitogen/helpers.py b/ansible_mitogen/helpers.py index 270f3244..b107d000 100644 --- a/ansible_mitogen/helpers.py +++ b/ansible_mitogen/helpers.py @@ -51,6 +51,9 @@ LOG = logging.getLogger(__name__) #: Caching of fetched file data. _file_cache = {} +#: Caching of compiled new-style module bytecode. +_bytecode_cache = {} + #: Mapping of job_id<->result dict _result_by_job_id = {} @@ -84,6 +87,13 @@ def get_file(context, path): return _file_cache[path] +def get_bytecode(context, path): + if path not in _bytecode_cache: + source = get_file(context, path) + _bytecode_cache[path] = compile(source, path, 'exec') + return _bytecode_cache[path] + + def run_module(kwargs): """ Set up the process environment in preparation for running an Ansible diff --git a/ansible_mitogen/planner.py b/ansible_mitogen/planner.py index b5f1bca0..ae632ab9 100644 --- a/ansible_mitogen/planner.py +++ b/ansible_mitogen/planner.py @@ -241,7 +241,7 @@ class WantJsonPlanner(ScriptPlanner): return 'WANT_JSON' in invocation.module_source -class NewStylePlanner(Planner): +class NewStylePlanner(ScriptPlanner): """ The Ansiballz framework differs from module replacer in that it uses real Python imports of things in ansible/module_utils instead of merely @@ -252,30 +252,6 @@ class NewStylePlanner(Planner): def detect(self, invocation): return 'from ansible.module_utils.' in invocation.module_source - def get_command_module_name(self, module_name): - """ - Given the name of an Ansible command module, return its canonical - module path within the ansible. - - :param module_name: - "shell" - :return: - "ansible.modules.commands.shell" - """ - path = module_loader.find_plugin(module_name, '') - relpath = os.path.relpath(path, os.path.dirname(ansible.__file__)) - root, _ = os.path.splitext(relpath) - return 'ansible.' + root.replace('/', '.') - - def plan(self, invocation): - return { - 'runner_name': self.runner_name, - 'module': invocation.module_name, - 'mod_name': self.get_command_module_name(invocation.module_name), - 'args': invocation.module_args, - 'env': invocation.env, - } - class ReplacerPlanner(NewStylePlanner): runner_name = 'ReplacerRunner' diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py index ed421cca..d6b3d48d 100644 --- a/ansible_mitogen/runner.py +++ b/ansible_mitogen/runner.py @@ -36,10 +36,13 @@ how to build arguments for it, preseed related data, etc. """ from __future__ import absolute_import +import cStringIO import json import logging import os +import sys import tempfile +import types import ansible_mitogen.helpers # TODO: circular import @@ -131,124 +134,37 @@ class TemporaryEnvironment(object): os.environ.update(self.original) -class NewStyleModuleExit(Exception): - """ - Capture the result of a call to `.exit_json()` or `.fail_json()` by a - native Ansible module. - """ - def __init__(self, ansible_module, **kwargs): - ansible_module.add_path_info(kwargs) - kwargs.setdefault('invocation', { - 'module_args': ansible_module.params - }) - ansible_module.do_cleanup_files() - self.dct = ansible.module_utils.basic.remove_values( - kwargs, - ansible_module.no_log_values, - ) - - -class NewStyleMethodOverrides(object): - @staticmethod - def exit_json(self, **kwargs): - """ - Raise exit_json() output as the `.dct` attribute of a - :class:`NewStyleModuleExit` exception`. - """ - kwargs.setdefault('changed', False) - raise NewStyleModuleExit(self, **kwargs) - - @staticmethod - def fail_json(self, **kwargs): - """ - Raise fail_json() output as the `.dct` attribute of a - :class:`NewStyleModuleExit` exception`. - """ - kwargs.setdefault('failed', True) - raise NewStyleModuleExit(self, **kwargs) - - klass = ansible.module_utils.basic.AnsibleModule - - def __init__(self): - self._original_exit_json = self.klass.exit_json - self._original_fail_json = self.klass.fail_json - self.klass.exit_json = self.exit_json - self.klass.fail_json = self.fail_json +class TemporaryArgv(object): + def __init__(self, argv): + self.original = sys.argv[:] + sys.argv[:] = argv def revert(self): - """ - Restore prior state. - """ - self.klass.exit_json = self._original_exit_json - self.klass.fail_json = self._original_fail_json + sys.argv[:] = self.original -class NewStyleModuleArguments(object): +class NewStyleStdio(object): """ Patch ansible.module_utils.basic argument globals. """ def __init__(self, args): - self.original = ansible.module_utils.basic._ANSIBLE_ARGS + self.original_stdout = sys.stdout + self.original_stderr = sys.stderr + self.original_stdin = sys.stdin + sys.stdout = cStringIO.StringIO() + sys.stderr = cStringIO.StringIO() ansible.module_utils.basic._ANSIBLE_ARGS = json.dumps({ 'ANSIBLE_MODULE_ARGS': args }) + sys.stdin = cStringIO.StringIO( + ansible.module_utils.basic._ANSIBLE_ARGS + ) def revert(self): - """ - Restore prior state. - """ - ansible.module_utils.basic._ANSIBLE_ARGS = self.original - - -class NewStyleRunner(Runner): - """ - Execute a new-style Ansible module, where Module Replacer-related tricks - aren't required. - """ - def __init__(self, mod_name, **kwargs): - super(NewStyleRunner, self).__init__(**kwargs) - self.mod_name = mod_name - - def setup(self): - super(NewStyleRunner, self).setup() - self._overrides = NewStyleMethodOverrides() - self._args = NewStyleModuleArguments(self.args) - - def revert(self): - super(NewStyleRunner, self).revert() - self._args.revert() - self._overrides.revert() - - def _fixup__default(self, mod): - pass - - def _fixup__yum_repository(self, mod): - # https://github.com/dw/mitogen/issues/154 - mod.YumRepo.repofile = mod.configparser.RawConfigParser() - - def _run(self): - fixup = getattr(self, '_fixup__' + self.module, self._fixup__default) - try: - mod = __import__(self.mod_name, {}, {}, ['']) - fixup(mod) - # Ansible modules begin execution on import. Thus the above - # __import__ will cause either Exit or ModuleError to be raised. If - # we reach the line below, the module did not execute and must - # already have been imported for a previous invocation, so we need - # to invoke main explicitly. - mod.main() - except NewStyleModuleExit, e: - return { - 'rc': 0, - 'stdout': json.dumps(e.dct), - 'stderr': '', - } - - return { - 'rc': 1, - 'stdout': '', - 'stderr': 'ansible_mitogen: module did not exit normally.', - } + sys.stdout = self.original_stdout + sys.stderr = self.original_stderr + sys.stdin = self.original_stdin + ansible.module_utils.basic._ANSIBLE_ARGS = '{}' class ProgramRunner(Runner): @@ -385,6 +301,47 @@ class ScriptRunner(ProgramRunner): return '\n'.join(new) +class NewStyleRunner(ScriptRunner): + """ + Execute a new-style Ansible module, where Module Replacer-related tricks + aren't required. + """ + def setup(self): + super(NewStyleRunner, self).setup() + self._stdio = NewStyleStdio(self.args) + self._argv = TemporaryArgv([self.path]) + + def revert(self): + super(NewStyleRunner, self).revert() + self._stdio.revert() + + def _get_bytecode(self): + """ + Fetch the module binary from the master if necessary. + """ + return ansible_mitogen.helpers.get_bytecode( + context=self.service_context, + path=self.path, + ) + + def _run(self): + bytecode = self._get_bytecode() + mod = types.ModuleType('__main__') + d = vars(mod) + e = None + + try: + exec bytecode in d, d + except SystemExit, e: + pass + + return { + 'rc': e[0] if e else 2, + 'stdout': sys.stdout.getvalue(), + 'stderr': sys.stderr.getvalue(), + } + + class JsonArgsFileRunner(ArgsFileRunner, ScriptRunner): JSON_ARGS = '<>' @@ -411,4 +368,4 @@ class OldStyleRunner(ArgsFileRunner, ScriptRunner): return ' '.join( '%s=%s' % (key, shlex_quote(str(self.args[key]))) for key in self.args - ) + ) + ' ' # Bug-for-bug :(