from __future__ import absolute_import # Let's not allow PyLint complain about string substitution # pylint: disable=W1321,E1321 # Import python libs import logging # Import Salt libs import salt.returners # Import third party libs try: import psycopg2 import psycopg2.extras HAS_POSTGRES = True except ImportError: HAS_POSTGRES = False LOG = logging.getLogger(__name__) def __virtual__(): if not HAS_POSTGRES: return False, 'Could not import saltresource module; psycopg2 is not installed.' return 'saltresource' def _get_options(ret=None): ''' Get the postgres options from salt. ''' attrs = {'host': 'host', 'user': 'user', 'passwd': 'passwd', 'db': 'db', 'port': 'port'} _options = salt.returners.get_returner_options('returner.postgres_graph_db', ret, attrs, __salt__=__salt__, __opts__=__opts__) return _options def _get_conn(ret=None): ''' Return a postgres connection. ''' _options = _get_options(ret) host = _options.get('host') user = _options.get('user') passwd = _options.get('passwd') datab = _options.get('db') port = _options.get('port') 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 graph_data(*args, **kwargs): ''' Returns graph data for visualization app CLI Examples: .. code-block:: bash salt '*' saltresource.graph_data ''' conn = _get_conn() cur_dict = conn.cursor(cursor_factory=psycopg2.extras.DictCursor) cur_dict.execute('SELECT host, service, status FROM salt_resources') resources_db = [dict(res) for res in cur_dict] db_dict = {} for resource in resources_db: host = resource.get('host') service = '.'.join(resource.get('service').split('.')[:2]) status = resource.get('status') if db_dict.get(host, None): if db_dict[host].get(service, None): service_data = db_dict[host][service] service_data.append(status) else: db_dict[host][service] = [status] else: db_dict[host] = {service: []} graph = [] for host, services in db_dict.items(): for service, statuses in services.items(): status = 'unknown' if 'failed' in statuses: status = 'failed' elif 'success' in statuses and not ('failed' in statuses or 'unknown' in statuses): status = 'success' datum = {'host': host, 'service': service, 'status': status} graph.append(datum) _close_conn(conn) return {'graph': graph} def host_data(host, **kwargs): ''' Returns data describing single host CLI Examples: .. code-block:: bash salt-call saltresource.host_data '' ''' conn = _get_conn() cur_dict = conn.cursor(cursor_factory=psycopg2.extras.DictCursor) sql = 'SELECT host, service, resource_id, last_ret, status FROM salt_resources WHERE host=%s' cur_dict.execute(sql, (host,)) resources_db = [dict(res) for res in cur_dict] db_dict = {} for resource in resources_db: host = resource.get('host') service = '.'.join(resource.get('service').split('.')[:2]) status = resource.get('status') if db_dict.get(host, None): if db_dict[host].get(service, None): service_data = db_dict[host][service] service_data.append(status) else: db_dict[host][service] = [status] else: db_dict[host] = {service: []} graph = [] for host, services in db_dict.items(): for service, statuses in services.items(): status = 'unknown' if 'failed' in statuses: status = 'failed' elif 'success' in statuses and not ('failed' in statuses or 'unknown' in statuses): status = 'success' resources = [{'service': r.get('service', ''), 'resource_id': r.get('resource_id', ''), 'last_ret': r.get('last_ret', None), 'status': r.get('status', '')} for r in resources_db if r.get('service', '').startswith(service)] datum = {'host': host, 'service': service, 'status': status, 'resources': resources} graph.append(datum) _close_conn(conn) return {'graph': graph} def sync_db(*args, **kwargs): conn = _get_conn() cur = conn.cursor() resources_sql = ''' CREATE TABLE IF NOT EXISTS salt_resources ( id varchar(255) NOT NULL UNIQUE, resource_id varchar(255) NOT NULL, host varchar(255) NOT NULL, service varchar(255) NOT NULL, module varchar(50) NOT NULL, fun varchar(50) NOT NULL, status varchar(50) NOT NULL, options json NULL, last_ret text NULL, alter_time TIMESTAMP WITH TIME ZONE DEFAULT now() ); ''' cur.execute(resources_sql) conn.commit() resources_meta_sql = ''' CREATE TABLE IF NOT EXISTS salt_resources_meta ( id varchar(255) NOT NULL UNIQUE, options json NULL, alter_time TIMESTAMP WITH TIME ZONE DEFAULT now() ); ''' cur.execute(resources_meta_sql) _close_conn(conn) return True def flush_db(*args, **kwargs): conn = _get_conn() cur = conn.cursor() result = True resources_sql = 'DELETE FROM salt_resources' try: cur.execute(resources_sql) conn.commit() except Exception as e: LOG.warning(repr(e)) result = False resources_meta_sql = 'DELETE FROM salt_resources_meta' try: cur.execute(resources_meta_sql) _close_conn(conn) except Exception as e: LOG.warning(repr(e)) result = False return result def destroy_db(*args, **kwargs): conn = _get_conn() cur = conn.cursor() resources_sql = 'DROP TABLE IF EXISTS salt_resources;' cur.execute(resources_sql) conn.commit() resources_meta_sql = 'DROP TABLE IF EXISTS salt_resources_meta;' cur.execute(resources_meta_sql) _close_conn(conn) return True