Merge pull request #11124 from sivel/jit-connect

Use a decorator to perform just in time connection

lib/ansible/executor/task_executor.py

@@ -210,7 +210,6 @@ class TaskExecutor:
         # get the connection and the handler for this execution
         self._connection = self._get_connection(variables)
         self._connection.set_host_overrides(host=self._host)
-        self._connection._connect()
 
         self._handler = self._get_action_handler(connection=self._connection, templar=templar)
 

lib/ansible/plugins/connections/__init__.py

@@ -22,6 +22,7 @@ __metaclass__ = type
 
 from abc import ABCMeta, abstractmethod, abstractproperty
 
+from functools import wraps
 from six import with_metaclass
 
 from ansible import constants as C
@@ -32,7 +33,16 @@ from ansible.errors import AnsibleError
 #        which may want to output display/logs too
 from ansible.utils.display import Display
 
-__all__ = ['ConnectionBase']
+__all__ = ['ConnectionBase', 'ensure_connect']
+
+
+def ensure_connect(func):
+    @wraps(func)
+    def wrapped(self, *args, **kwargs):
+        self._connect()
+        return func(self, *args, **kwargs)
+    return wrapped
+
 
 class ConnectionBase(with_metaclass(ABCMeta, object)):
     '''
@@ -82,16 +92,19 @@ class ConnectionBase(with_metaclass(ABCMeta, object)):
         """Connect to the host we've been initialized with"""
         pass
 
+    @ensure_connect
     @abstractmethod
     def exec_command(self, cmd, tmp_path, executable=None, in_data=None):
         """Run a command on the remote host"""
         pass
 
+    @ensure_connect
     @abstractmethod
     def put_file(self, in_path, out_path):
         """Transfer a file from local to remote"""
         pass
 
+    @ensure_connect
     @abstractmethod
     def fetch_file(self, in_path, out_path):
         """Fetch a file from remote to local"""

lib/ansible/plugins/connections/local.py

@@ -49,6 +49,8 @@ class Connection(ConnectionBase):
     def exec_command(self, cmd, tmp_path, executable='/bin/sh', in_data=None):
         ''' run a command on the local host '''
 
+        super(Connection, self).exec_command(cmd, tmp_path, executable=executable, in_data=in_data)
+
         debug("in local.exec_command()")
         # su requires to be run from a terminal, and therefore isn't supported here (yet?)
         #if self._connection_info.su:
@@ -108,6 +110,8 @@ class Connection(ConnectionBase):
     def put_file(self, in_path, out_path):
         ''' transfer a file from local to local '''
 
+        super(Connection, self).put_file(in_path, out_path)
+
         #vvv("PUT {0} TO {1}".format(in_path, out_path), host=self.host)
         self._display.vvv("{0} PUT {1} TO {2}".format(self._connection_info.remote_addr, in_path, out_path))
         if not os.path.exists(in_path):
@@ -123,6 +127,9 @@ class Connection(ConnectionBase):
 
     def fetch_file(self, in_path, out_path):
         ''' fetch a file from local to local -- for copatibility '''
+
+        super(Connection, self).fetch_file(in_path, out_path)
+
         #vvv("FETCH {0} TO {1}".format(in_path, out_path), host=self.host)
         self._display.vvv("{0} FETCH {1} TO {2}".format(self._connection_info.remote_addr, in_path, out_path))
         self.put_file(in_path, out_path)

lib/ansible/plugins/connections/paramiko_ssh.py

@@ -61,6 +61,7 @@ with warnings.catch_warnings():
     except ImportError:
         pass
 
+
 class MyAddPolicy(object):
     """
     Based on AutoAddPolicy in paramiko so we can determine when keys are added
@@ -191,6 +192,8 @@ class Connection(ConnectionBase):
     def exec_command(self, cmd, tmp_path, executable='/bin/sh', in_data=None):
         ''' run a command on the remote host '''
 
+        super(Connection, self).exec_command(cmd, tmp_path, executable=executable, in_data=in_data)
+
         if in_data:
             raise AnsibleError("Internal Error: this module does not support optimized module pipelining")
 
@@ -251,6 +254,8 @@ class Connection(ConnectionBase):
     def put_file(self, in_path, out_path):
         ''' transfer a file from local to remote '''
 
+        super(Connection, self).put_file(in_path, out_path)
+
         self._display.vvv("PUT %s TO %s" % (in_path, out_path), host=self._connection_info.remote_addr)
 
         if not os.path.exists(in_path):
@@ -272,12 +277,14 @@ class Connection(ConnectionBase):
         if cache_key in SFTP_CONNECTION_CACHE:
             return SFTP_CONNECTION_CACHE[cache_key]
         else:
-            result = SFTP_CONNECTION_CACHE[cache_key] = self.connect().ssh.open_sftp()
+            result = SFTP_CONNECTION_CACHE[cache_key] = self._connect().ssh.open_sftp()
             return result
 
     def fetch_file(self, in_path, out_path):
         ''' save a remote file to the specified path '''
 
+        super(Connection, self).fetch_file(in_path, out_path)
+
         self._display.vvv("FETCH %s TO %s" % (in_path, out_path), host=self._connection_info.remote_addr)
 
         try:

lib/ansible/plugins/connections/ssh.py

@@ -36,6 +36,7 @@ from ansible import constants as C
 from ansible.errors import AnsibleError, AnsibleConnectionFailure, AnsibleFileNotFound
 from ansible.plugins.connections import ConnectionBase
 
+
 class Connection(ConnectionBase):
     ''' ssh based connections '''
 
@@ -272,6 +273,8 @@ class Connection(ConnectionBase):
     def exec_command(self, cmd, tmp_path, executable='/bin/sh', in_data=None):
         ''' run a command on the remote host '''
 
+        super(Connection, self).exec_command(cmd, tmp_path, executable=executable, in_data=in_data)
+
         ssh_cmd = self._password_cmd()
         ssh_cmd += ("ssh", "-C")
         if not in_data:
@@ -392,6 +395,9 @@ class Connection(ConnectionBase):
 
     def put_file(self, in_path, out_path):
         ''' transfer a file from local to remote '''
+
+        super(Connection, self).put_file(in_path, out_path)
+
         self._display.vvv("PUT {0} TO {1}".format(in_path, out_path), host=self._connection_info.remote_addr)
         if not os.path.exists(in_path):
             raise AnsibleFileNotFound("file or module does not exist: {0}".format(in_path))
@@ -427,6 +433,9 @@ class Connection(ConnectionBase):
 
     def fetch_file(self, in_path, out_path):
         ''' fetch a file from remote to local '''
+
+        super(Connection, self).fetch_file(in_path, out_path)
+
         self._display.vvv("FETCH {0} TO {1}".format(in_path, out_path), host=self._connection_info.remote_addr)
         cmd = self._password_cmd()
 

lib/ansible/plugins/connections/winrm.py

@@ -46,6 +46,7 @@ from ansible.plugins.connections import ConnectionBase
 from ansible.plugins import shell_loader
 from ansible.utils.path import makedirs_safe
 
+
 class Connection(ConnectionBase):
     '''WinRM connections over HTTP/HTTPS.'''
 
@@ -152,6 +153,7 @@ class Connection(ConnectionBase):
         return self
 
     def exec_command(self, cmd, tmp_path, executable='/bin/sh', in_data=None):
+        super(Connection, self).exec_command(cmd, tmp_path, executable=executable, in_data,in_data)
 
         cmd = cmd.encode('utf-8')
         cmd_parts = shlex.split(cmd, posix=False)
@@ -173,6 +175,8 @@ class Connection(ConnectionBase):
         return (result.status_code, '', result.std_out.encode('utf-8'), result.std_err.encode('utf-8'))
 
     def put_file(self, in_path, out_path):
+        super(Connection, self).put_file(in_path, out_path)
+
         self._display.vvv("PUT %s TO %s" % (in_path, out_path), host=self._connection_info.remote_addr)
         if not os.path.exists(in_path):
             raise AnsibleFileNotFound("file or module does not exist: %s" % in_path)
@@ -211,6 +215,8 @@ class Connection(ConnectionBase):
                     raise AnsibleError("failed to transfer file to %s" % out_path)
 
     def fetch_file(self, in_path, out_path):
+        super(Connection, self).fetch_file(in_path, out_path)
+
         out_path = out_path.replace('\\', '/')
         self._display.vvv("FETCH %s TO %s" % (in_path, out_path), host=self._connection_info.remote_addr)
         buffer_size = 2**19 # 0.5MB chunks