diff --git a/ansible_mitogen/connection/mitogen.py b/ansible_mitogen/connection/mitogen.py index a00826bd..e6a3486a 100644 --- a/ansible_mitogen/connection/mitogen.py +++ b/ansible_mitogen/connection/mitogen.py @@ -43,6 +43,11 @@ class Connection(ansible.plugins.connection.ConnectionBase): become_methods = ['sudo'] transport = 'mitogen' + def __init__(self, play_context, new_stdin, original_transport): + self.original_transport = original_transport + self.transport = original_transport + super(Connection, self).__init__(play_context, new_stdin) + @property def connected(self): return self.router is not None @@ -80,19 +85,17 @@ class Connection(ansible.plugins.connection.ConnectionBase): path = os.environ['LISTENER_SOCKET_PATH'] self.router, self.parent = mitogen.unix.connect(path) - if self._play_context.connection == 'local': - host = self._connect_local() - else: - host = self._connect_ssh() - - if not self._play_context.become: - self.context = host + if self.original_transport == 'local': + self.context = self._connect_local() else: - self.context = self._connect_sudo(via=host) + self.host = self._connect_ssh() + if not self._play_context.become: + self.context = self.host + else: + self.context = self._connect_sudo(via=self.host) def call_async(self, func, *args, **kwargs): self._connect() - print[func, args, kwargs] return self.context.call_async(func, *args, **kwargs) def call(self, func, *args, **kwargs): diff --git a/ansible_mitogen/strategy/mitogen.py b/ansible_mitogen/strategy/mitogen.py index c1f85d05..d0363d6e 100644 --- a/ansible_mitogen/strategy/mitogen.py +++ b/ansible_mitogen/strategy/mitogen.py @@ -63,6 +63,19 @@ action_loader__get = ansible.plugins.action_loader.get ansible.plugins.action_loader.get = wrap_action_loader__get +def wrap_connection_loader__get(name, play_context, new_stdin): + """ + """ + kwargs = {} + if name in ('ssh', 'local'): + kwargs['original_transport'] = name + name = 'mitogen' + return connection_loader__get(name, play_context, new_stdin, **kwargs) + +connection_loader__get = ansible.plugins.connection_loader.get +ansible.plugins.connection_loader.get = wrap_connection_loader__get + + class ContextProxyService(mitogen.service.Service): """ Implement a service accessible from worker processes connecting back into @@ -111,10 +124,7 @@ class StrategyModule(ansible.plugins.strategy.linear.StrategyModule): os.environ['LISTENER_SOCKET_PATH'] = self.listener.path self.service = ContextProxyService(self.router) - mitogen.utils.log_to_file()#level='DEBUG', io=False) - - if play_context.connection == 'ssh': - play_context.connection = 'mitogen' + #mitogen.utils.log_to_file(level='DEBUG', io=False) import threading th = threading.Thread(target=self.service.run)