Salt Engine for graph data processing

Change-Id: I08ab88c3cdbe3914973ce87039222f29e1dff824

_engines/saltgraph.py

+# -*- coding: utf-8 -*-
+'''
+Saltgraph engine for catching returns of state runs, parsing them 
+and passing them to flat database of latest Salt resource runs.
+'''
+
+# Import python libs
+from __future__ import absolute_import
+import datetime
+import json
+import logging
+
+# Import salt libs
+import salt.utils.event
+
+# Import third party libs
+try:
+    import psycopg2
+    import psycopg2.extras
+    HAS_POSTGRES = True
+except ImportError:
+    HAS_POSTGRES = False
+
+__virtualname__ = 'saltgraph'
+log = logging.getLogger(__name__)
+
+
+def __virtual__():
+    if not HAS_POSTGRES:
+        return False, 'Could not import saltgraph engine. python-psycopg2 is not installed.'
+    return __virtualname__
+
+
+def _get_conn(options={}):
+    '''
+    Return a postgres connection.
+    '''
+    host = options.get('host', '127.0.0.1')
+    user = options.get('user', 'salt')
+    passwd = options.get('passwd', 'salt')
+    datab = options.get('db', 'salt')
+    port = options.get('port', 5432)
+
+    return psycopg2.connect(
+            host=host,
+            user=user,
+            password=passwd,
+            database=datab,
+            port=port)
+
+
+def _close_conn(conn):
+    '''
+    Close the Postgres connection
+    '''
+    conn.commit()
+    conn.close()
+
+
+def _get_lowstate_data(options={}):
+    '''
+    TODO: document this method
+    '''
+    conn = _get_conn(options)
+    cur = conn.cursor()
+
+    try:
+        # you can only do this on Salt Masters minion
+        lowstate_req = __salt__['saltutil.cmd']('*', 'state.show_lowstate', **{'timeout': 15, 'concurrent': True, 'queue': True})
+    except:
+        lowstate_req = {}
+
+    for minion, lowstate_ret in lowstate_req.items():
+        if lowstate_ret.get('retcode') != 0:
+            continue
+        for resource in lowstate_ret.get('ret', []):
+            low_sql = '''INSERT INTO salt_resources
+                         (id, resource_id, host, service, module, fun, status)
+                         VALUES (%s, %s, %s, %s, %s, %s, %s)
+                         ON CONFLICT (id) DO UPDATE
+                           SET resource_id = excluded.resource_id,
+                               host = excluded.host,
+                               service = excluded.service,
+                               module = excluded.module,
+                               fun = excluded.fun,
+                               alter_time = current_timestamp'''
+
+            rid = "%s|%s" % (minion, resource.get('__id__'))
+            cur.execute(
+                low_sql, (
+                    rid,
+                    resource.get('__id__'),
+                    minion,
+                    resource.get('__sls__'),
+                    resource.get('state') if 'state' in resource else resource.get('module'),
+                    resource.get('fun'),
+                    'unknown'
+                )
+            )
+            conn.commit()
+
+    if lowstate_req:
+        meta_sql = '''INSERT INTO salt_resources_meta
+                      (id, options)
+                      VALUES (%s, %s)
+                      ON CONFLICT (id) DO UPDATE
+                        SET options = excluded.options,
+                            alter_time = current_timestamp'''
+
+        cur.execute(
+            meta_sql, (
+                'lowstate_data',
+                '{}'
+            )
+        )
+    _close_conn(conn)
+
+
+def _up_to_date(options={}):
+    '''
+    TODO: document this method
+    '''
+    conn = _get_conn(options)
+    cur = conn.cursor()
+    #cur_dict = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+
+    ret = False
+
+    # if lowstate data are older than 1 day, refresh them
+    cur.execute('SELECT alter_time FROM salt_resources_meta WHERE id = %s', ('lowstate_data',))
+    alter_time = cur.fetchone()
+
+    if alter_time:
+        now = datetime.datetime.utcnow()
+        day = datetime.timedelta(days=1)
+        time_diff = now - alter_time[0].replace(tzinfo=None)
+        if time_diff < day:
+            ret = True
+    else:
+        skip = False
+
+    _close_conn(conn)
+
+    return ret
+
+
+def _update_resources(event, options):
+    '''
+    TODO: document this method
+    '''
+    conn = _get_conn(options)
+    cur = conn.cursor()
+
+    cur.execute('SELECT id FROM salt_resources')
+    resources_db = [res[0] for res in cur.fetchall()]
+    resources = event.get('return', {}).values()
+
+    for resource in resources:
+        rid = '%s|%s' % (event.get('id'), resource.get('__id__'))
+        if rid in resources_db:
+            status = 'unknown'
+            if resource.get('result', None) is not None:
+                status = 'success' if resource.get('result') else 'failed'
+
+            resource_sql = '''UPDATE salt_resources SET (status, last_ret, alter_time) = (%s, %s, current_timestamp)
+                                WHERE id = %s'''
+
+            cur.execute(
+                resource_sql, (
+                    status,
+                    repr(resource),
+                    rid
+                )
+            )
+
+            conn.commit()
+
+    _close_conn(conn)
+
+
+def start(host='salt', user='salt', password='salt', database='salt', port=5432, **kwargs):
+    '''
+    Listen to events and parse Salt state returns
+    '''
+    if __opts__['__role'] == 'master':
+        event_bus = salt.utils.event.get_master_event(
+                __opts__,
+                __opts__['sock_dir'],
+                listen=True)
+    else:
+        event_bus = salt.utils.event.get_event(
+            'minion',
+            transport=__opts__['transport'],
+            opts=__opts__,
+            sock_dir=__opts__['sock_dir'],
+            listen=True)
+        log.debug('Saltgraph engine started')
+
+    while True:
+        event = event_bus.get_event()
+        supported_funcs = ['state.sls', 'state.apply', 'state.highstate']
+        if event and event.get('fun', None) in supported_funcs:
+            test = 'test=true' in [arg.lower() for arg in event.get('fun_args', [])]
+            if not test:
+                options = {
+                    'host': host,
+                    'user': user,
+                    'passwd': password,
+                    'db': database,
+                    'port': port
+                }
+                is_reclass = [arg for arg in event.get('fun_args', []) if arg.startswith('reclass')]
+                if is_reclass or not _up_to_date(options):
+                    _get_lowstate_data(options)
+
+                _update_resources(event, options)
+