<s> import pymongo <EOL> from hm import config <EOL> from hm . model import host , load_balancer <EOL> class MongoDBStorage ( object ) : <EOL> hosts_collection = "<STR_LIT>" <EOL> lb_collection = "<STR_LIT>" <EOL> def __init__ ( self , conf = None ) : <EOL> self . config = conf <EOL> self . mongo_uri = config . get_config ( '<STR_LIT>' , None , conf ) <EOL> if not self . mongo_uri : <EOL> self . mongo_uri = config . get_config ( '<STR_LIT>' , '<STR_LIT>' , conf ) <EOL> client = pymongo . MongoClient ( self . mongo_uri ) <EOL> try : <EOL> self . db = client . get_default_database ( ) <EOL> self . mongo_database = self . db . name <EOL> except pymongo . errors . ConfigurationError : <EOL> self . mongo_database = config . get_config ( '<STR_LIT>' , '<STR_LIT>' , conf ) <EOL> self . db = client [ self . mongo_database ] <EOL> def store_host ( self , h ) : <EOL> self . _hosts_collection ( ) . insert ( h . to_json ( ) ) <EOL> def remove_host ( self , id ) : <EOL> self . _hosts_collection ( ) . remove ( { '<STR_LIT>' : id } ) <EOL> def find_host ( self , id ) : <EOL> h_data = self . _hosts_collection ( ) . find_one ( { '<STR_LIT>' : id } ) <EOL> return host . Host . from_dict ( h_data , conf = self . config ) <EOL> def list_hosts ( self , filters ) : <EOL> host_data_list = self . _hosts_collection ( ) . find ( filters ) or [ ] <EOL> return [ host . Host . from_dict ( h_data , conf = self . config ) for h_data in host_data_list ] <EOL> def store_load_balancer ( self , lb ) : <EOL> self . _lb_collection ( ) . insert ( lb . to_json ( ) ) <EOL> def remove_load_balancer ( self , name ) : <EOL> self . _lb_collection ( ) . remove ( { '<STR_LIT>' : name } ) <EOL> def find_load_balancer ( self , name ) : <EOL> lb_data = self . _lb_collection ( ) . find_one ( name ) <EOL> return load_balancer . LoadBalancer . from_dict ( lb_data , conf = self . config ) <EOL> def list_load_balancers ( self , filters ) : <EOL> lb_data_list = self . _lb_collection ( ) . find ( filters ) or [ ] <EOL> return [ load_balancer . LoadBalancer . from_dict ( lb_data , conf = self . config ) for lb_data in lb_data_list ] <EOL> def add_host_to_load_balancer ( self , name , h ) : <EOL> self . _lb_collection ( ) . update ( { '<STR_LIT>' : name } , { '<STR_LIT>' : { '<STR_LIT>' : h . to_json ( ) } } ) <EOL> def remove_host_from_load_balancer ( self , name , h ) : <EOL> self . _lb_collection ( ) . update ( { '<STR_LIT>' : name } , { '<STR_LIT>' : { '<STR_LIT>' : { '<STR_LIT>' : h . id } } } ) <EOL> def _hosts_collection ( self ) : <EOL> return self . db [ self . hosts_collection ] <EOL> def _lb_collection ( self ) : <EOL> return self . db [ self . lb_collection ] </s>
