<s> import copy <EOL> import importlib <EOL> import ssl as ssl_lib <EOL> import six <EOL> import mongoengine <EOL> from st2common import log as logging <EOL> from st2common . util import isotime <EOL> from st2common . models . db import stormbase <EOL> from st2common . models . utils . profiling import log_query_and_profile_data_for_queryset <EOL> LOG = logging . getLogger ( __name__ ) <EOL> MODEL_MODULE_NAMES = [ <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> ] <EOL> def get_model_classes ( ) : <EOL> """<STR_LIT>""" <EOL> result = [ ] <EOL> for module_name in MODEL_MODULE_NAMES : <EOL> module = importlib . import_module ( module_name ) <EOL> model_classes = getattr ( module , '<STR_LIT>' , [ ] ) <EOL> result . extend ( model_classes ) <EOL> return result <EOL> def db_setup ( db_name , db_host , db_port , username = None , password = None , <EOL> ensure_indexes = True , ssl = False , ssl_keyfile = None , ssl_certfile = None , <EOL> ssl_cert_reqs = None , ssl_ca_certs = None , ssl_match_hostname = True ) : <EOL> LOG . info ( '<STR_LIT>' , <EOL> db_name , db_host , db_port , str ( username ) ) <EOL> ssl_kwargs = _get_ssl_kwargs ( ssl = ssl , ssl_keyfile = ssl_keyfile , ssl_certfile = ssl_certfile , <EOL> ssl_cert_reqs = ssl_cert_reqs , ssl_ca_certs = ssl_ca_certs , <EOL> ssl_match_hostname = ssl_match_hostname ) <EOL> connection = mongoengine . connection . connect ( db_name , host = db_host , <EOL> port = db_port , tz_aware = True , <EOL> username = username , password = password , <EOL> ** ssl_kwargs ) <EOL> if ensure_indexes : <EOL> db_ensure_indexes ( ) <EOL> return connection <EOL> def db_ensure_indexes ( ) : <EOL> """<STR_LIT>""" <EOL> LOG . debug ( '<STR_LIT>' ) <EOL> model_classes = get_model_classes ( ) <EOL> for cls in model_classes : <EOL> LOG . debug ( '<STR_LIT>' % ( cls . __name__ ) ) <EOL> cls . ensure_indexes ( ) <EOL> def db_teardown ( ) : <EOL> mongoengine . connection . disconnect ( ) <EOL> def _get_ssl_kwargs ( ssl = False , ssl_keyfile = None , ssl_certfile = None , ssl_cert_reqs = None , <EOL> ssl_ca_certs = None , ssl_match_hostname = True ) : <EOL> ssl_kwargs = { <EOL> '<STR_LIT>' : ssl , <EOL> } <EOL> if ssl_keyfile : <EOL> ssl_kwargs [ '<STR_LIT>' ] = True <EOL> ssl_kwargs [ '<STR_LIT>' ] = ssl_keyfile <EOL> if ssl_certfile : <EOL> ssl_kwargs [ '<STR_LIT>' ] = True <EOL> ssl_kwargs [ '<STR_LIT>' ] = ssl_certfile <EOL> if ssl_cert_reqs : <EOL> if ssl_cert_reqs is '<STR_LIT:none>' : <EOL> ssl_cert_reqs = ssl_lib . CERT_NONE <EOL> elif ssl_cert_reqs is '<STR_LIT>' : <EOL> ssl_cert_reqs = ssl_lib . CERT_OPTIONAL <EOL> elif ssl_cert_reqs is '<STR_LIT>' : <EOL> ssl_cert_reqs = ssl_lib . CERT_REQUIRED <EOL> ssl_kwargs [ '<STR_LIT>' ] = ssl_cert_reqs <EOL> if ssl_ca_certs : <EOL> ssl_kwargs [ '<STR_LIT>' ] = True <EOL> ssl_kwargs [ '<STR_LIT>' ] = ssl_ca_certs <EOL> if ssl_kwargs . get ( '<STR_LIT>' , False ) : <EOL> ssl_kwargs [ '<STR_LIT>' ] = ssl_match_hostname <EOL> return ssl_kwargs <EOL> class MongoDBAccess ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , model ) : <EOL> self . model = model <EOL> def get_by_name ( self , value ) : <EOL> return self . get ( name = value , raise_exception = True ) <EOL> def get_by_id ( self , value ) : <EOL> return self . get ( id = value , raise_exception = True ) <EOL> def get_by_ref ( self , value ) : <EOL> return self . get ( ref = value , raise_exception = True ) <EOL> def get ( self , exclude_fields = None , * args , ** kwargs ) : <EOL> raise_exception = kwargs . pop ( '<STR_LIT>' , False ) <EOL> instances = self . model . objects ( ** kwargs ) <EOL> if exclude_fields : <EOL> instances = instances . exclude ( * exclude_fields ) <EOL> instance = instances [ <NUM_LIT:0> ] if instances else None <EOL> log_query_and_profile_data_for_queryset ( queryset = instances ) <EOL> if not instance and raise_exception : <EOL> raise ValueError ( '<STR_LIT>' % ( self . model . __name__ , kwargs ) ) <EOL> return instance <EOL> def get_all ( self , * args , ** kwargs ) : <EOL> return self . query ( * args , ** kwargs ) <EOL> def count ( self , * args , ** kwargs ) : <EOL> result = self . model . objects ( ** kwargs ) . count ( ) <EOL> log_query_and_profile_data_for_queryset ( queryset = result ) <EOL> return result <EOL> def query ( self , offset = <NUM_LIT:0> , limit = None , order_by = None , exclude_fields = None , <EOL> ** filters ) : <EOL> order_by = order_by or [ ] <EOL> exclude_fields = exclude_fields or [ ] <EOL> eop = offset + int ( limit ) if limit else None <EOL> filters , order_by = self . _process_datetime_range_filters ( filters = filters , order_by = order_by ) <EOL> filters = self . _process_null_filters ( filters = filters ) <EOL> result = self . model . objects ( ** filters ) <EOL> if exclude_fields : <EOL> result = result . exclude ( * exclude_fields ) <EOL> result = result . order_by ( * order_by ) <EOL> result = result [ offset : eop ] <EOL> log_query_and_profile_data_for_queryset ( queryset = result ) <EOL> return result <EOL> def distinct ( self , * args , ** kwargs ) : <EOL> field = kwargs . pop ( '<STR_LIT>' ) <EOL> result = self . model . objects ( ** kwargs ) . distinct ( field ) <EOL> log_query_and_profile_data_for_queryset ( queryset = result ) <EOL> return result <EOL> def aggregate ( self , * args , ** kwargs ) : <EOL> return self . model . objects ( ** kwargs ) . _collection . aggregate ( * args , ** kwargs ) <EOL> def insert ( self , instance ) : <EOL> instance = self . model . objects . insert ( instance ) <EOL> return self . _undo_dict_field_escape ( instance ) <EOL> def add_or_update ( self , instance ) : <EOL> instance . save ( ) <EOL> return self . _undo_dict_field_escape ( instance ) <EOL> def update ( self , instance , ** kwargs ) : <EOL> return instance . update ( ** kwargs ) <EOL> def delete ( self , instance ) : <EOL> return instance . delete ( ) <EOL> def delete_by_query ( self , ** query ) : <EOL> qs = self . model . objects . filter ( ** query ) <EOL> qs . delete ( ) <EOL> log_query_and_profile_data_for_queryset ( queryset = qs ) <EOL> return None <EOL> def _undo_dict_field_escape ( self , instance ) : <EOL> for attr , field in instance . _fields . iteritems ( ) : <EOL> if isinstance ( field , stormbase . EscapedDictField ) : <EOL> value = getattr ( instance , attr ) <EOL> setattr ( instance , attr , field . to_python ( value ) ) <EOL> return instance <EOL> def _process_null_filters ( self , filters ) : <EOL> result = copy . deepcopy ( filters ) <EOL> null_filters = { k : v for k , v in six . iteritems ( filters ) <EOL> if v is None or ( type ( v ) in [ str , unicode ] and str ( v . lower ( ) ) == '<STR_LIT:null>' ) } <EOL> for key in null_filters . keys ( ) : <EOL> result [ '<STR_LIT>' % ( key ) ] = False <EOL> del result [ key ] <EOL> return result <EOL> def _process_datetime_range_filters ( self , filters , order_by = None ) : <EOL> ranges = { k : v for k , v in filters . iteritems ( ) <EOL> if type ( v ) in [ str , unicode ] and '<STR_LIT:..>' in v } <EOL> order_by_list = copy . deepcopy ( order_by ) if order_by else [ ] <EOL> for k , v in ranges . iteritems ( ) : <EOL> values = v . split ( '<STR_LIT:..>' ) <EOL> dt1 = isotime . parse ( values [ <NUM_LIT:0> ] ) <EOL> dt2 = isotime . parse ( values [ <NUM_LIT:1> ] ) <EOL> k__gte = '<STR_LIT>' % k <EOL> k__lte = '<STR_LIT>' % k <EOL> if dt1 < dt2 : <EOL> query = { k__gte : dt1 , k__lte : dt2 } <EOL> sort_key , reverse_sort_key = k , '<STR_LIT:->' + k <EOL> else : <EOL> query = { k__gte : dt2 , k__lte : dt1 } <EOL> sort_key , reverse_sort_key = '<STR_LIT:->' + k , k <EOL> del filters [ k ] <EOL> filters . update ( query ) <EOL> if reverse_sort_key in order_by_list : <EOL> idx = order_by_list . index ( reverse_sort_key ) <EOL> order_by_list . pop ( idx ) <EOL> order_by_list . insert ( idx , sort_key ) <EOL> elif sort_key not in order_by_list : <EOL> order_by_list = [ sort_key ] + order_by_list <EOL> return filters , order_by_list </s>
