ansible: handle local connections and synchronize module too.

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):

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)