<s> '''<STR_LIT>''' <EOL> from __future__ import absolute_import <EOL> from contextlib import contextmanager <EOL> import sys <EOL> import time <EOL> import logging <EOL> import salt . returners <EOL> import salt . utils . jid <EOL> import salt . exceptions <EOL> try : <EOL> import psycopg2 <EOL> import psycopg2 . extras <EOL> HAS_PG = True <EOL> except ImportError : <EOL> HAS_PG = False <EOL> log = logging . getLogger ( __name__ ) <EOL> __virtualname__ = '<STR_LIT>' <EOL> def __virtual__ ( ) : <EOL> if not HAS_PG : <EOL> return False <EOL> return True <EOL> def _get_options ( ret = None ) : <EOL> '''<STR_LIT>''' <EOL> defaults = { '<STR_LIT:host>' : '<STR_LIT:localhost>' , <EOL> '<STR_LIT:user>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:port>' : <NUM_LIT> } <EOL> attrs = { '<STR_LIT:host>' : '<STR_LIT:host>' , <EOL> '<STR_LIT:user>' : '<STR_LIT:user>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:port>' : '<STR_LIT:port>' } <EOL> _options = salt . returners . get_returner_options ( '<STR_LIT>' . format ( __virtualname__ ) , <EOL> ret , <EOL> attrs , <EOL> __salt__ = __salt__ , <EOL> __opts__ = __opts__ , <EOL> defaults = defaults ) <EOL> if '<STR_LIT:port>' in _options : <EOL> _options [ '<STR_LIT:port>' ] = int ( _options [ '<STR_LIT:port>' ] ) <EOL> return _options <EOL> @ contextmanager <EOL> def _get_serv ( ret = None , commit = False ) : <EOL> '''<STR_LIT>''' <EOL> _options = _get_options ( ret ) <EOL> try : <EOL> ssl_options = { } <EOL> if _options . get ( '<STR_LIT>' ) : <EOL> ssl_options [ '<STR_LIT>' ] = _options . get ( '<STR_LIT>' ) <EOL> if _options . get ( '<STR_LIT>' ) : <EOL> ssl_options [ '<STR_LIT>' ] = _options . get ( '<STR_LIT>' ) <EOL> if _options . get ( '<STR_LIT>' ) : <EOL> ssl_options [ '<STR_LIT:key>' ] = _options . get ( '<STR_LIT>' ) <EOL> conn = psycopg2 . connect ( host = _options . get ( '<STR_LIT:host>' ) , <EOL> user = _options . get ( '<STR_LIT:user>' ) , <EOL> password = _options . get ( '<STR_LIT>' ) , <EOL> database = _options . get ( '<STR_LIT>' ) , <EOL> port = _options . get ( '<STR_LIT:port>' ) ) <EOL> except psycopg2 . OperationalError as exc : <EOL> raise salt . exceptions . SaltMasterError ( '<STR_LIT>' . format ( exc = exc ) ) <EOL> cursor = conn . cursor ( ) <EOL> try : <EOL> yield cursor <EOL> except psycopg2 . DatabaseError as err : <EOL> error = err . args <EOL> sys . stderr . write ( str ( error ) ) <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> raise err <EOL> else : <EOL> if commit : <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> else : <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> finally : <EOL> conn . close ( ) <EOL> def returner ( ret ) : <EOL> '''<STR_LIT>''' <EOL> try : <EOL> with _get_serv ( ret , commit = True ) as cur : <EOL> sql = '''<STR_LIT>''' <EOL> cur . execute ( sql , ( ret [ '<STR_LIT>' ] , ret [ '<STR_LIT>' ] , <EOL> psycopg2 . extras . Json ( ret [ '<STR_LIT>' ] ) , <EOL> ret [ '<STR_LIT:id>' ] , <EOL> ret . get ( '<STR_LIT:success>' , False ) , <EOL> psycopg2 . extras . Json ( ret ) , <EOL> time . strftime ( '<STR_LIT>' , time . localtime ( ) ) ) ) <EOL> except salt . exceptions . SaltMasterError : <EOL> log . critical ( '<STR_LIT>' ) <EOL> def event_return ( events ) : <EOL> '''<STR_LIT>''' <EOL> with _get_serv ( events , commit = True ) as cur : <EOL> for event in events : <EOL> tag = event . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> data = event . get ( '<STR_LIT:data>' , '<STR_LIT>' ) <EOL> sql = '''<STR_LIT>''' <EOL> cur . execute ( sql , ( tag , psycopg2 . extras . Json ( data ) , <EOL> __opts__ [ '<STR_LIT:id>' ] , time . strftime ( '<STR_LIT>' , time . localtime ( ) ) ) ) <EOL> def save_load ( jid , load ) : <EOL> '''<STR_LIT>''' <EOL> with _get_serv ( commit = True ) as cur : <EOL> sql = '''<STR_LIT>''' <EOL> try : <EOL> cur . execute ( sql , ( jid , psycopg2 . extras . Json ( load ) ) ) <EOL> except psycopg2 . IntegrityError : <EOL> pass <EOL> def save_minions ( jid , minions ) : <EOL> '''<STR_LIT>''' <EOL> pass <EOL> def get_load ( jid ) : <EOL> '''<STR_LIT>''' <EOL> with _get_serv ( ret = None , commit = True ) as cur : <EOL> sql = '''<STR_LIT>''' <EOL> cur . execute ( sql , ( jid , ) ) <EOL> data = cur . fetchone ( ) <EOL> if data : <EOL> return data [ <NUM_LIT:0> ] <EOL> return { } <EOL> def get_jid ( jid ) : <EOL> '''<STR_LIT>''' <EOL> with _get_serv ( ret = None , commit = True ) as cur : <EOL> sql = '''<STR_LIT>''' <EOL> cur . execute ( sql , ( jid , ) ) <EOL> data = cur . fetchall ( ) <EOL> ret = { } <EOL> if data : <EOL> for minion , full_ret in data : <EOL> ret [ minion ] = full_ret <EOL> return ret <EOL> def get_fun ( fun ) : <EOL> '''<STR_LIT>''' <EOL> with _get_serv ( ret = None , commit = True ) as cur : <EOL> sql = '''<STR_LIT>''' <EOL> cur . execute ( sql , ( fun , ) ) <EOL> data = cur . fetchall ( ) <EOL> ret = { } <EOL> if data : <EOL> for minion , _ , full_ret in data : <EOL> ret [ minion ] = full_ret <EOL> return ret <EOL> def get_jids ( ) : <EOL> '''<STR_LIT>''' <EOL> with _get_serv ( ret = None , commit = True ) as cur : <EOL> sql = '''<STR_LIT>''' <EOL> cur . execute ( sql ) <EOL> data = cur . fetchall ( ) <EOL> ret = { } <EOL> for jid , load in data : <EOL> ret [ jid ] = salt . utils . jid . format_jid_instance ( jid , load ) <EOL> return ret <EOL> def get_minions ( ) : <EOL> '''<STR_LIT>''' <EOL> with _get_serv ( ret = None , commit = True ) as cur : <EOL> sql = '''<STR_LIT>''' <EOL> cur . execute ( sql ) <EOL> data = cur . fetchall ( ) <EOL> ret = [ ] <EOL> for minion in data : <EOL> ret . append ( minion [ <NUM_LIT:0> ] ) <EOL> return ret <EOL> def prep_jid ( nocache = False , passed_jid = None ) : <EOL> '''<STR_LIT>''' <EOL> return passed_jid if passed_jid is not None else salt . utils . jid . gen_jid ( ) </s>
