<s> """<STR_LIT>""" <EOL> import pymongo <EOL> import gridfs <EOL> import hashlib <EOL> import datetime <EOL> import bson <EOL> import time <EOL> class DataStore ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , uri = '<STR_LIT>' , database = '<STR_LIT>' , worker_cap = <NUM_LIT:0> , samples_cap = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> self . sample_collection = '<STR_LIT>' <EOL> self . worker_cap = worker_cap <EOL> self . samples_cap = samples_cap <EOL> self . database_name = database <EOL> self . uri = '<STR_LIT>' + uri + '<STR_LIT:/>' + self . database_name <EOL> self . mongo = pymongo . MongoClient ( self . uri , use_greenlets = True ) <EOL> self . database = self . mongo . get_default_database ( ) <EOL> self . gridfs_handle = gridfs . GridFS ( self . database ) <EOL> self . last_ops_run = time . time ( ) <EOL> self . periodic_ops ( ) <EOL> print '<STR_LIT>' % ( self . uri , self . database_name ) <EOL> def get_uri ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . uri <EOL> def store_sample ( self , sample_bytes , filename , type_tag ) : <EOL> """<STR_LIT>""" <EOL> if len ( filename ) > <NUM_LIT:1000> : <EOL> print '<STR_LIT>' % ( sample_bytes [ : <NUM_LIT:100> ] , filename [ : <NUM_LIT:100> ] ) <EOL> exit ( <NUM_LIT:1> ) <EOL> sample_info = { } <EOL> sample_info [ '<STR_LIT>' ] = hashlib . md5 ( sample_bytes ) . hexdigest ( ) <EOL> if self . has_sample ( sample_info [ '<STR_LIT>' ] ) : <EOL> return sample_info [ '<STR_LIT>' ] <EOL> self . periodic_ops ( ) <EOL> self . expire_data ( ) <EOL> sample_info [ '<STR_LIT:filename>' ] = filename <EOL> sample_info [ '<STR_LIT>' ] = len ( sample_bytes ) <EOL> sample_info [ '<STR_LIT>' ] = datetime . datetime . utcnow ( ) <EOL> sample_info [ '<STR_LIT>' ] = type_tag <EOL> import random <EOL> sample_info [ '<STR_LIT>' ] = random . choice ( [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> sample_info [ '<STR_LIT>' ] = self . gridfs_handle . put ( sample_bytes ) <EOL> self . database [ self . sample_collection ] . insert ( sample_info ) <EOL> print '<STR_LIT>' % ( self . sample_storage_size ( ) , self . samples_cap ) <EOL> return sample_info [ '<STR_LIT>' ] <EOL> def sample_storage_size ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> coll_stats = self . database . command ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> sample_storage_size = coll_stats [ '<STR_LIT:size>' ] / <NUM_LIT> / <NUM_LIT> <EOL> return sample_storage_size <EOL> except pymongo . errors . OperationFailure : <EOL> return <NUM_LIT:0> <EOL> def expire_data ( self ) : <EOL> """<STR_LIT>""" <EOL> while self . sample_storage_size ( ) > self . samples_cap : <EOL> record = self . database [ self . sample_collection ] . find ( ) . sort ( '<STR_LIT>' , pymongo . ASCENDING ) . limit ( <NUM_LIT:1> ) [ <NUM_LIT:0> ] <EOL> self . remove_sample ( record [ '<STR_LIT>' ] ) <EOL> def remove_sample ( self , md5 ) : <EOL> """<STR_LIT>""" <EOL> record = self . database [ self . sample_collection ] . find_one ( { '<STR_LIT>' : md5 } ) <EOL> if not record : <EOL> return <EOL> print '<STR_LIT>' % ( record [ '<STR_LIT>' ] , record [ '<STR_LIT>' ] / <NUM_LIT> / <NUM_LIT> ) <EOL> self . database [ self . sample_collection ] . remove ( { '<STR_LIT>' : record [ '<STR_LIT>' ] } ) <EOL> self . gridfs_handle . delete ( record [ '<STR_LIT>' ] ) <EOL> print '<STR_LIT>' % ( self . sample_storage_size ( ) , self . samples_cap ) <EOL> def clean_for_serialization ( self , data ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( data , dict ) : <EOL> for k in data . keys ( ) : <EOL> if ( k . startswith ( '<STR_LIT>' ) ) : <EOL> del data [ k ] <EOL> elif isinstance ( data [ k ] , bson . objectid . ObjectId ) : <EOL> del data [ k ] <EOL> elif isinstance ( data [ k ] , datetime . datetime ) : <EOL> data [ k ] = data [ k ] . isoformat ( ) + '<STR_LIT>' <EOL> elif isinstance ( data [ k ] , dict ) : <EOL> data [ k ] = self . clean_for_serialization ( data [ k ] ) <EOL> elif isinstance ( data [ k ] , list ) : <EOL> data [ k ] = [ self . clean_for_serialization ( item ) for item in data [ k ] ] <EOL> return data <EOL> def clean_for_storage ( self , data ) : <EOL> """<STR_LIT>""" <EOL> data = self . data_to_unicode ( data ) <EOL> if isinstance ( data , dict ) : <EOL> for k in dict ( data ) . keys ( ) : <EOL> if k == '<STR_LIT>' : <EOL> del data [ k ] <EOL> continue <EOL> if '<STR_LIT:.>' in k : <EOL> new_k = k . replace ( '<STR_LIT:.>' , '<STR_LIT:_>' ) <EOL> data [ new_k ] = data [ k ] <EOL> del data [ k ] <EOL> k = new_k <EOL> if isinstance ( data [ k ] , dict ) : <EOL> data [ k ] = self . clean_for_storage ( data [ k ] ) <EOL> elif isinstance ( data [ k ] , list ) : <EOL> data [ k ] = [ self . clean_for_storage ( item ) for item in data [ k ] ] <EOL> return data <EOL> def get_full_md5 ( self , partial_md5 , collection ) : <EOL> """<STR_LIT>""" <EOL> print '<STR_LIT>' <EOL> starts_with = '<STR_LIT>' % partial_md5 <EOL> sample_info = self . database [ collection ] . find_one ( { '<STR_LIT>' : { '<STR_LIT>' : starts_with } } , { '<STR_LIT>' : <NUM_LIT:1> } ) <EOL> return sample_info [ '<STR_LIT>' ] if sample_info else None <EOL> def get_sample ( self , md5 ) : <EOL> """<STR_LIT>""" <EOL> if len ( md5 ) < <NUM_LIT:32> : <EOL> md5 = self . get_full_md5 ( md5 , self . sample_collection ) <EOL> sample_info = self . database [ self . sample_collection ] . find_one ( { '<STR_LIT>' : md5 } ) <EOL> if not sample_info : <EOL> return None <EOL> try : <EOL> grid_fs_id = sample_info [ '<STR_LIT>' ] <EOL> sample_info = self . clean_for_serialization ( sample_info ) <EOL> sample_info . update ( { '<STR_LIT>' : self . gridfs_handle . get ( grid_fs_id ) . read ( ) } ) <EOL> return sample_info <EOL> except gridfs . errors . CorruptGridFile : <EOL> self . database [ self . sample_collection ] . update ( { '<STR_LIT>' : md5 } , { '<STR_LIT>' : None } ) <EOL> return None <EOL> def get_sample_window ( self , type_tag , size = <NUM_LIT:10> ) : <EOL> """<STR_LIT>""" <EOL> size = size * <NUM_LIT> * <NUM_LIT> <EOL> cursor = self . database [ self . sample_collection ] . find ( { '<STR_LIT>' : type_tag } , <EOL> { '<STR_LIT>' : <NUM_LIT:1> , '<STR_LIT>' : <NUM_LIT:1> } ) . sort ( '<STR_LIT>' , pymongo . DESCENDING ) <EOL> total_size = <NUM_LIT:0> <EOL> md5_list = [ ] <EOL> for item in cursor : <EOL> if total_size > size : <EOL> return md5_list <EOL> md5_list . append ( item [ '<STR_LIT>' ] ) <EOL> total_size += item [ '<STR_LIT>' ] <EOL> return md5_list <EOL> def has_sample ( self , md5 ) : <EOL> """<STR_LIT>""" <EOL> sample = self . get_sample ( md5 ) <EOL> return True if sample else False <EOL> def _list_samples ( self , predicate = None ) : <EOL> """<STR_LIT>""" <EOL> cursor = self . database [ self . sample_collection ] . find ( predicate , { '<STR_LIT>' : <NUM_LIT:0> , '<STR_LIT>' : <NUM_LIT:1> } ) <EOL> return [ item [ '<STR_LIT>' ] for item in cursor ] <EOL> def tag_match ( self , tags = None ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' not in self . database . collection_names ( ) : <EOL> print '<STR_LIT>' <EOL> return None <EOL> if not tags : <EOL> cursor = self . database [ '<STR_LIT>' ] . find ( { } , { '<STR_LIT>' : <NUM_LIT:0> , '<STR_LIT>' : <NUM_LIT:1> } ) <EOL> else : <EOL> cursor = self . database [ '<STR_LIT>' ] . find ( { '<STR_LIT>' : { '<STR_LIT>' : tags } } , { '<STR_LIT>' : <NUM_LIT:0> , '<STR_LIT>' : <NUM_LIT:1> } ) <EOL> tag_md5s = set ( [ item [ '<STR_LIT>' ] for item in cursor ] ) <EOL> sample_md5s = set ( item [ '<STR_LIT>' ] for item in self . database [ '<STR_LIT>' ] . find ( { } , { '<STR_LIT>' : <NUM_LIT:0> , '<STR_LIT>' : <NUM_LIT:1> } ) ) <EOL> return list ( tag_md5s . intersection ( sample_md5s ) ) <EOL> def tags_all ( self ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' not in self . database . collection_names ( ) : <EOL> print '<STR_LIT>' <EOL> return None <EOL> cursor = self . database [ '<STR_LIT>' ] . find ( { } , { '<STR_LIT>' : <NUM_LIT:0> , '<STR_LIT>' : <NUM_LIT:1> , '<STR_LIT>' : <NUM_LIT:1> } ) <EOL> return [ item for item in cursor ] <EOL> def store_work_results ( self , results , collection , md5 ) : <EOL> """<STR_LIT>""" <EOL> results [ '<STR_LIT>' ] = md5 <EOL> results [ '<STR_LIT>' ] = datetime . datetime . utcnow ( ) <EOL> if '<STR_LIT>' not in results : <EOL> results [ '<STR_LIT>' ] = results [ '<STR_LIT>' ] <EOL> try : <EOL> self . database [ collection ] . update ( { '<STR_LIT>' : md5 } , self . clean_for_storage ( results ) , True ) <EOL> except pymongo . errors . OperationFailure : <EOL> print '<STR_LIT>' <EOL> print '<STR_LIT>' % ( collection , md5 ) <EOL> def get_work_results ( self , collection , md5 ) : <EOL> """<STR_LIT>""" <EOL> return self . database [ collection ] . find_one ( { '<STR_LIT>' : md5 } ) <EOL> def all_sample_md5s ( self , type_tag = None ) : <EOL> """<STR_LIT>""" <EOL> if type_tag : <EOL> cursor = self . database [ self . sample_collection ] . find ( { '<STR_LIT>' : type_tag } , { '<STR_LIT>' : <NUM_LIT:1> , '<STR_LIT>' : <NUM_LIT:0> } ) <EOL> else : <EOL> cursor = self . database [ self . sample_collection ] . find ( { } , { '<STR_LIT>' : <NUM_LIT:1> , '<STR_LIT>' : <NUM_LIT:0> } ) <EOL> return [ match . values ( ) [ <NUM_LIT:0> ] for match in cursor ] <EOL> def clear_worker_output ( self ) : <EOL> """<STR_LIT>""" <EOL> print '<STR_LIT>' <EOL> all_c = self . database . collection_names ( ) <EOL> try : <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( self . sample_collection ) <EOL> except ValueError : <EOL> print '<STR_LIT>' <EOL> for collection in all_c : <EOL> self . database . drop_collection ( collection ) <EOL> def clear_db ( self ) : <EOL> """<STR_LIT>""" <EOL> print '<STR_LIT>' <EOL> self . mongo . drop_database ( self . database_name ) <EOL> def periodic_ops ( self ) : <EOL> """<STR_LIT>""" <EOL> if ( time . time ( ) - self . last_ops_run ) < <NUM_LIT:30> : <EOL> return <EOL> try : <EOL> self . last_ops_run = time . time ( ) <EOL> print '<STR_LIT>' <EOL> all_c = self . database . collection_names ( ) <EOL> try : <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( '<STR_LIT:info>' ) <EOL> all_c . remove ( '<STR_LIT>' ) <EOL> all_c . remove ( self . sample_collection ) <EOL> except ValueError : <EOL> print '<STR_LIT>' <EOL> if self . worker_cap : <EOL> size = self . worker_cap * pow ( <NUM_LIT> , <NUM_LIT:2> ) <EOL> for collection in all_c : <EOL> self . database . command ( '<STR_LIT>' , collection , size = size ) <EOL> for collection in all_c : <EOL> self . database [ collection ] . ensure_index ( '<STR_LIT>' ) <EOL> self . database [ self . sample_collection ] . create_index ( '<STR_LIT>' ) <EOL> self . database [ '<STR_LIT>' ] . create_index ( '<STR_LIT>' ) <EOL> except pymongo . errors . AutoReconnect as e : <EOL> print '<STR_LIT>' % e <EOL> return <EOL> except Exception as e : <EOL> print '<STR_LIT>' % e <EOL> return <EOL> def to_unicode ( self , s ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( s , unicode ) : <EOL> return s <EOL> if isinstance ( s , str ) : <EOL> return unicode ( s , errors = '<STR_LIT:ignore>' ) <EOL> return s <EOL> def data_to_unicode ( self , data ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( data , dict ) : <EOL> return { self . to_unicode ( k ) : self . to_unicode ( v ) for k , v in data . iteritems ( ) } <EOL> if isinstance ( data , list ) : <EOL> return [ self . to_unicode ( l ) for l in data ] <EOL> else : <EOL> return self . to_unicode ( data ) </s>
