added session proxy

Allow to override session object for some command and all its children.
Proxy encapsulates main session object and provides limited access to
its uris and connections.

4 files changed, 69 insertions, 14 deletions

diff --git a/lmi/scripts/common/command/base.py b/lmi/scripts/common/command/base.py
index 2998f63..d1f6cc7 100644
--- a/lmi/scripts/common/command/base.py
+++ b/lmi/scripts/common/command/base.py
@@ -173,6 +173,20 @@ class LmiBaseCommand(object):
             options = self.parent.format_options
         return options
 
+    @property
+    def session(self):
+        """
+        :returns: Session object. Session for command and all of its children
+            may be overriden with a call to :py:meth:`set_session_proxy`.
+        :rtype: :py:class:`lmi.scripts.common.session.Session`
+        """
+        proxy = getattr(self, '_session_proxy', None)
+        if proxy:
+            return proxy
+        if self.parent is not None:
+            return self.parent.session
+        return self.app.session
+
     def get_cmd_name_parts(self, all_parts=False, demand_own_usage=True,
             for_docopt=False):
         """
@@ -281,3 +295,16 @@ class LmiBaseCommand(object):
         :rtype: integer
         """
         raise NotImplementedError("run method must be overriden in subclass")
+
+    def set_session_proxy(self, session):
+        """
+        Allows to override session object. This is useful for especially for
+        conditional commands (subclasses of
+        :py:class:`~lmi.scripts.common.command.LmiSelectCommand`) that devide
+        connections to groups satisfying particular expression. These groups
+        are turned into session proxies containing just a subset of connections
+        in global session object.
+
+        :param session: Session object.
+        """
+        self._session_proxy = session
diff --git a/lmi/scripts/common/command/checkresult.py b/lmi/scripts/common/command/checkresult.py
index 08b3b63..27e9549 100644
--- a/lmi/scripts/common/command/checkresult.py
+++ b/lmi/scripts/common/command/checkresult.py
@@ -134,7 +134,7 @@ class LmiCheckResult(LmiSessionCommand):
         pass
 
     def process_session_results(self, session, results):
-        if len(self.app.session) > 1:
+        if len(self.session) > 1:
             LOG().debug('Successful runs: %d\n',
                     len([r for r in results.values() if r[0]]))
         LmiSessionCommand.process_session_results(self, session, results)
diff --git a/lmi/scripts/common/command/session.py b/lmi/scripts/common/command/session.py
index a617539..9d88bbe 100644
--- a/lmi/scripts/common/command/session.py
+++ b/lmi/scripts/common/command/session.py
@@ -127,7 +127,7 @@ class LmiSessionCommand(LmiEndPointCommand):
             successful invocation or an exception.
         """
         if success:
-            if len(self.app.session) > 1:
+            if len(self.session) > 1:
                 self.formatter.print_host(hostname)
             self.produce_output(result)
 
@@ -187,7 +187,7 @@ class LmiSessionCommand(LmiEndPointCommand):
         desired and modifies connection accordingly.
 
         :param connection: Connection to single host.
-        :type connection: :py:class:`lmi.shell.LMIConnection` 
+        :type connection: :py:class:`lmi.shell.LMIConnection`
         :param list args: Arguments handed over to associated function.
         :param dictionary kwargs: Keyword arguments handed over to associated
             function.
@@ -202,5 +202,5 @@ class LmiSessionCommand(LmiEndPointCommand):
         return self.execute(connection, *args, **kwargs)
 
     def run_with_args(self, args, kwargs):
-        return self.process_session(self.app.session, args, kwargs)
+        return self.process_session(self.session, args, kwargs)
 
diff --git a/lmi/scripts/common/session.py b/lmi/scripts/common/session.py
index 381bae5..b063d0a 100644
--- a/lmi/scripts/common/session.py
+++ b/lmi/scripts/common/session.py
@@ -35,6 +35,7 @@ from collections import defaultdict
 
 from lmi.scripts.common import errors
 from lmi.scripts.common import get_logger
+from lmi.scripts.common.util import FilteredDict
 from lmi.shell.LMIConnection import connect
 
 LOG = get_logger(__name__)
@@ -76,26 +77,29 @@ class Session(object):
             ``None`` if connection can not be made.
         """
         if self._connections[hostname] is None:
-            self._connections[hostname] = self._connect(
-                    hostname, interactive=True)
+            try:
+                self._connections[hostname] = self._connect(
+                        hostname, interactive=True)
+            except Exception as exc:
+                LOG().error('Failed to make a connection to "%s": %s',
+                        hostname, exc)
         return self._connections[hostname]
 
     def __len__(self):
         """ Get the number of hostnames in session. """
         return len(self._connections)
 
+    def __contains__(self, uri):
+        return uri in self._connections
+
     def __iter__(self):
         """ Yields connection objects. """
         successful_connections = 0
         for hostname in self._connections:
-            try:
-                connection = self[hostname]
-                if connection is not None:
-                    yield connection
-                    successful_connections += 1
-            except Exception as exc:
-                LOG().error('Failed to make a connection to "%s": %s',
-                        hostname, exc)
+            connection = self[hostname]
+            if connection is not None:
+                yield connection
+                successful_connections += 1
         if successful_connections == 0:
             raise errors.LmiNoConnections('No successful connection made.')
 
@@ -160,3 +164,27 @@ class Session(object):
         """
         return [h for h, c in self._connections.items() if c is None]
 
+class SessionProxy(Session):
+    """
+    Behaves like a session. But it just encapsulates other session object and
+    provides access to a subset of its items.
+
+    :param session: Session object or even another session proxy.
+    :param list uris: Subset of uris in encapsulated session object.
+    """
+
+    def __init__(self, session, uris):
+        uris = set(uris)
+        if not all(isinstance(uri, basestring) for uri in uris):
+            raise ValueError("uris must be iterable of uris")
+        for uri in uris:
+            if not uri in session:
+                raise ValueError('uri "%s" needs to belong to given session'
+                        % uri)
+        Session.__init__(self, session._app, uris, session._credentials,
+                session._same_credentials)
+        self._origin = session
+        self._connections = FilteredDict(uris, session._connections)
+        # let the credentials propagage to original session
+        self._credentials = session._credentials
+