mysql-dynamic-inventory.py

from ansible.errors import AnsibleError
from ansible.plugins.inventory import BaseInventoryPlugin
from ansible.utils.display import Display
import pymysql

display = Display()

DOCUMENTATION = r"""
    name: mysql-dynamic-inventory
    plugin_type: inventory
    short_description: Returns Ansible inventory from an SQL query
    description: Returns Ansible inventory from an SQL query
    options:
      plugin:
          description: Name of the plugin
          required: true
          choices: ['mysql-dynamic-inventory']
      db_host:
        description: Database host
        required: true
      db_user:
        description: Database user
        required: true
      db_pass:
        description: Database password
        required: true
      db_name:
        description: Database name
        required: true
      db_query:
        description: Database query
        required: true
"""


class InventoryModule(BaseInventoryPlugin):
    NAME = "mysql-dynamic-inventory"

    def verify_file(self, path):
        """Ensures that the given file is valid for this plugin"""
        valid = False
        if super(InventoryModule, self).verify_file(path):
            if path.endswith(("mysql.yaml", "mysql.yml")):
                valid = True
        return valid

    def parse(self, inventory, loader, path, cache=True):
        """Parses the inventory file"""
        super(InventoryModule, self).parse(inventory, loader, path, cache)
        self._read_config_data(path)

        self.templar.available_variables = self._vars

        for option in self._options:
            try:
                if self.get_option(option):
                    templated_option = self.templar.template(self.get_option(option))
                    self.set_option(option, templated_option)
            except Exception as e:
                raise AnsibleError(
                    f"Error processing templating for {option}: {str(e)}"
                )

        self._fetch_hosts()

    def _fetch_hosts(self):
        """Fetches hosts from the database and populates the inventory"""
        group_init = {}
        try:
            connection = pymysql.connect(
                host=self.get_option("db_host"),
                user=self.get_option("db_user"),
                password=self.get_option("db_pass"),
                database=self.get_option("db_name"),
                cursorclass=pymysql.cursors.DictCursor,
            )
            with connection.cursor() as cursor:
                cursor.execute(self.get_option("db_query"))
                for row in cursor.fetchall():
                    group = row.get("inventory_group")
                    if group and group not in group_init:
                        self.inventory.add_group(group)
                        group_init[group] = True
                    hostname = row.get("inventory_hostname")
                    if hostname:
                        self.inventory.add_host(hostname, group=group)
                        # Dynamically set ansible host variables for each column returned
                        for key, value in row.items():
                            # Skip inventory_hostname and group as they're already used
                            if key not in ["inventory_hostname", "inventory_group"]:
                                self.inventory.set_variable(hostname, key, value)
        except Exception as e:
            raise AnsibleError(f"Database query failed: {e}")
        finally:
            if connection:
                connection.close()