Refactor lago ansible_hosts

lago/cmd.py

@@ -26,20 +26,15 @@
 import os
 import pkg_resources
 import sys
+from textwrap import dedent
 import warnings
 
-from collections import defaultdict
-
 import lago
 import lago.plugins
 import lago.plugins.cli
 import lago.templates
 from lago.config import config
-from lago import (
-    log_utils,
-    workdir as lago_workdir,
-    utils,
-)
+from lago import (log_utils, workdir as lago_workdir, utils, lago_ansible)
 from lago.utils import (in_prefix, with_logging, LagoUserException)
 
 LOGGER = logging.getLogger('cli')
@@ -489,47 +484,67 @@ def do_console(prefix, host, **kwargs):
 
 
 @lago.plugins.cli.cli_plugin(
-    help='Create Ansible host inventory of the environment'
+    help='Create Ansible host inventory of the environment',
+    description=dedent(
+        """
+    This method iterates through all the VMs and creates an Ansible
+    host inventory. For each vm it defines an IP address and a private key.
+
+    The default groups are based on the values which associated
+    with the following keys: 'vm-type', 'groups', 'vm-provider'.
+
+    The 'keys' parameter can be used to override the default groups,
+    for example to create a group which based on a 'service_provider',
+    --keys 'service_provider' should be added to this command.
+
+    Nested keys can be used also by specifying the path to the key,
+    for example '/disks/0/metadata/distro' will create a group based on the
+    distro of the os installed on disk at position 0 in the init file.
+    (we assume that numeric values in the path should be used as index for
+    list access).
+
+    If the value associated with the key is a list, a group will be created
+    for every item in that list (this is useful when you want to associate
+    a machine with more than one group).
+
+    The output of this command is printed to standard output.
+
+    Example of a possible output:
+
+    lago ansible_hosts -k 'vm-type'
+
+    [vm-type=ovirt-host]
+    lago-host1 ansible_host=1.2.3.4 ansible_ssh_private_key_file=/path/rsa
+    lago-host2 ansible_host=1.2.3.6 ansible_ssh_private_key_file=/path/rsa
+    [vm-type=ovirt-engine]
+    lago-engine ansible_host=1.2.3.5 ansible_ssh_private_key_file=/path/rsa
+
+    lago ansible_hosts -k 'disks/0/metadata/arch' 'groups'
+
+    [disks/0/metadata/arch=x86_64]
+    vm0-server ansible_host=1.2.3.4 ansible_ssh_private_key_file=/path/rsa
+    vm1-slave ansible_host=1.2.3.5 ansible_ssh_private_key_file=/path/rsa
+    vm2-slave ansible_host=1.2.3.6 ansible_ssh_private_key_file=/path/rsa
+    [groups=slaves]
+    vm1-slave ansible_host=1.2.3.5 ansible_ssh_private_key_file=/path/rsa
+    vm2-slave ansible_host=1.2.3.6 ansible_ssh_private_key_file=/path/rsa
+    [groups=servers]
+    vm0-server ansible_host=1.2.3.4 ansible_ssh_private_key_file=/path/rsa
+    """
+    )
+)
+@lago.plugins.cli.cli_plugin_add_argument(
+    '--keys',
+    '-k',
+    help='Path to the keys that should be used as groups',
+    default=['vm-type', 'groups', 'vm-provider'],
+    metavar='KEY',
+    nargs='*',
 )
 @in_lago_prefix
 @with_logging
-def do_generate_ansible_hosts(prefix, out_format, **kwargs):
-    """
-    Create Ansible host inventory of the environment
-    """
-    #
-    # This method iterates through all the VMs and creates an Ansible
-    # host inventory. It defines an IP address and private key
-    # for the machine in group named by its type, it is printed to
-    # standard output. The content looks for example like this:
-    #
-    # [ovirt-host]
-    # lago-host1 ansible_host=1.2.3.4 ansible_ssh_private_key_file=/path/rsa
-    # lago-host2 ansible_host=1.2.3.6 ansible_ssh_private_key_file=/path/rsa
-    # [ovirt-engine]
-    # lago-engine ansible_host=1.2.3.5 ansible_ssh_private_key_file=/path/rsa
-    #
-    inventory = defaultdict(list)
-    nets = prefix.get_nets()
-    for vm in prefix.virt_env.get_vms().values():
-        vm_mgmt_ip = [
-            nic.get('ip') for nic in vm.nics()
-            if nets[nic.get('net')].is_management()
-        ]
-        for ip in vm_mgmt_ip:
-            inventory[vm._spec['vm-type']].append(
-                "{name} ansible_host={ip} ansible_ssh_private_key_file={key}"
-                .format(
-                    ip=ip,
-                    key=prefix.virt_env.prefix.paths.ssh_id_rsa(),
-                    name=vm.name(),
-                )
-            )
-
-    for name, hosts in inventory.iteritems():
-        print('[{name}]'.format(name=name))
-        for host in sorted(hosts):
-            print(host)
+def do_generate_ansible_hosts(prefix, keys, **kwargs):
+    print(lago_ansible.LagoAnsible(prefix).get_inventory_str(keys))
 
 
 @lago.plugins.cli.cli_plugin(

lago/lago_ansible.py

@@ -0,0 +1,159 @@
+from collections import defaultdict
+import functools
+import logging
+import tempfile
+import contextlib
+
+from . import log_utils
+
+LOGGER = logging.getLogger(__name__)
+LogTask = functools.partial(log_utils.LogTask, logger=LOGGER)
+
+
+class LagoAnsible(object):
+    """
+    A class for handling Ansible related tasks
+
+    Attributes:
+        prefix (lago.prefix.Prefix): The prefix that this object wraps
+    """
+
+    def __init__(self, prefix):
+        """
+        Args:
+            prefix (lago.prefix.Prefix): A prefix to wrap
+        """
+        self.prefix = prefix
+
+    def get_inventory_str(self, keys=None):
+        """
+        Convert a dict generated by ansible.LagoAnsible.get_inventory
+        to an INI-like file.
+
+        Args:
+            keys (list of str): Path to the keys that will be used to
+                create groups.
+
+        Returns:
+            str: INI-like Ansible inventory
+        """
+        inventory = self.get_inventory(keys)
+        lines = []
+        for name, hosts in inventory.viewitems():
+            lines.append('[{name}]'.format(name=name))
+            for host in sorted(hosts):
+                lines.append(host)
+
+        return '\n'.join(lines)
+
+    def get_inventory(self, keys=None):
+        """
+        Create an Ansible inventory based on python dicts and lists.
+        The returned value is a dict in which every key represents a group
+        and every value is a list of entries for that group.
+
+        Args:
+            keys (list of str): Path to the keys that will be used to
+                create groups.
+
+        Returns:
+            dict: dict based Ansible inventory
+        """
+        inventory = defaultdict(list)
+        keys = keys or ['vm-type', 'groups', 'vm-provider']
+        vms = self.prefix.get_vms().values()
+
+        for vm in vms:
+            entry = self._generate_entry(vm)
+            vm_spec = vm.spec
+            for key in keys:
+                value = self.get_key(key, vm_spec)
+                if value is None:
+                    continue
+                if isinstance(value, list):
+                    for sub_value in value:
+                        inventory['{}={}'.format(key, sub_value)].append(entry)
+                else:
+                    inventory['{}={}'.format(key, value)].append(entry)
+
+        return inventory
+
+    def _generate_entry(self, vm):
+        """
+        Generate host entry for the given VM
+        Args:
+            vm (lago.plugins.vm.VMPlugin): The VM for which the entry
+                should be created for.
+
+        Returns:
+            str: An entry for vm
+        """
+        return \
+            '{name} ' \
+            'ansible_host={ip} ' \
+            'ansible_ssh_private_key_file={key}'.format(
+                name=vm.name(),
+                ip=vm.ip(),
+                key=self.prefix.paths.ssh_id_rsa()
+            )
+
+    @staticmethod
+    def get_key(key, data_structure):
+        """
+        Helper method for extracting values from a nested data structure.
+
+        Args:
+            key (str): The path to the vales (a series of keys and indexes
+                separated by '/')
+            data_structure (dict or list): The data structure from which the
+                value will be extracted.
+
+        Returns:
+            str: The values associated with key
+        """
+        if key == '/':
+            return data_structure
+
+        path = key.split('/')
+        # If the path start with '/', remove the empty string from the list
+        path[0] or path.pop(0)
+        current_value = data_structure
+        while path:
+            current_key = path.pop(0)
+            try:
+                current_key = int(current_key)
+            except ValueError:
+                pass
+
+            try:
+                current_value = current_value[current_key]
+            except (KeyError, IndexError):
+                LOGGER.debug('failed to extract path {}'.format(key))
+                return None
+
+        return current_value
+
+    @contextlib.contextmanager
+    def get_inventory_temp_file(self, keys=None):
+        """
+        Context manager which returns the inventory written on a tempfile.
+        The tempfile will be deleted as soon as this context manger ends.
+
+        Args:
+            keys (list of str): Path to the keys that will be used to
+                create groups.
+
+        Yields:
+            tempfile.NamedTemporaryFile: Temp file containing the inventory
+        """
+        temp_file = tempfile.NamedTemporaryFile(mode='r+t')
+        inventory = self.get_inventory_str(keys)
+        LOGGER.debug(
+            'Writing inventory to temp file {} \n{}'.
+            format(temp_file.name, inventory)
+        )
+        temp_file.write(inventory)
+        temp_file.flush()
+        temp_file.seek(0)
+        yield temp_file
+        temp_file.close()

lago/plugins/vm.py

@@ -419,6 +419,23 @@ def mgmt_name(self):
     def mgmt_net(self):
         return self.virt_env.get_net(name=self.mgmt_name)
 
+    @property
+    def vm_type(self):
+        return self._spec['vm-type']
+
+    @property
+    def groups(self):
+        """
+        Returns:
+            list of str: The names of the groups to which this vm belongs
+                (as specified in the init file)
+        """
+        groups = self._spec.get('groups', [])
+        if groups:
+            return groups[:]
+        else:
+            return groups
+
     def name(self):
         return str(self._spec['name'])