<s> import datetime <EOL> import json <EOL> from pkg_resources import iter_entry_points <EOL> import sqlalchemy as sql <EOL> from moai . utils import check_type <EOL> def get_database ( uri , config = None ) : <EOL> prefix = uri . split ( '<STR_LIT::>' ) [ <NUM_LIT:0> ] <EOL> for entry_point in iter_entry_points ( group = '<STR_LIT>' , name = prefix ) : <EOL> dbclass = entry_point . load ( ) <EOL> try : <EOL> return dbclass ( uri , config ) <EOL> except TypeError : <EOL> return dbclass ( uri ) <EOL> raise ValueError ( '<STR_LIT>' % prefix ) <EOL> class SQLDatabase ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , dburi = None ) : <EOL> self . _uri = dburi <EOL> self . _db = self . _connect ( ) <EOL> self . _records = self . _db . tables [ '<STR_LIT>' ] <EOL> self . _sets = self . _db . tables [ '<STR_LIT>' ] <EOL> self . _setrefs = self . _db . tables [ '<STR_LIT>' ] <EOL> self . _reset_cache ( ) <EOL> def _connect ( self ) : <EOL> dburi = self . _uri <EOL> if dburi is None : <EOL> dburi = '<STR_LIT>' <EOL> engine = sql . create_engine ( dburi ) <EOL> db = sql . MetaData ( engine ) <EOL> sql . Table ( '<STR_LIT>' , db , <EOL> sql . Column ( '<STR_LIT>' , sql . Unicode , primary_key = True ) , <EOL> sql . Column ( '<STR_LIT>' , sql . DateTime , index = True ) , <EOL> sql . Column ( '<STR_LIT>' , sql . Boolean ) , <EOL> sql . Column ( '<STR_LIT>' , sql . String ) ) <EOL> sql . Table ( '<STR_LIT>' , db , <EOL> sql . Column ( '<STR_LIT>' , sql . Unicode , primary_key = True ) , <EOL> sql . Column ( '<STR_LIT>' , sql . Boolean ) , <EOL> sql . Column ( '<STR_LIT:name>' , sql . Unicode ) , <EOL> sql . Column ( '<STR_LIT:description>' , sql . Unicode ) ) <EOL> sql . Table ( '<STR_LIT>' , db , <EOL> sql . Column ( '<STR_LIT>' , sql . Integer , <EOL> sql . ForeignKey ( '<STR_LIT>' ) , <EOL> index = True , primary_key = True ) , <EOL> sql . Column ( '<STR_LIT>' , sql . Integer , <EOL> sql . ForeignKey ( '<STR_LIT>' ) , <EOL> index = True , primary_key = True ) ) <EOL> db . create_all ( ) <EOL> return db <EOL> def flush ( self ) : <EOL> oai_ids = set ( ) <EOL> for row in sql . select ( [ self . _records . c . record_id ] ) . execute ( ) : <EOL> oai_ids . add ( row [ <NUM_LIT:0> ] ) <EOL> for row in sql . select ( [ self . _sets . c . set_id ] ) . execute ( ) : <EOL> oai_ids . add ( row [ <NUM_LIT:0> ] ) <EOL> deleted_records = [ ] <EOL> deleted_sets = [ ] <EOL> deleted_setrefs = [ ] <EOL> inserted_records = [ ] <EOL> inserted_sets = [ ] <EOL> inserted_setrefs = [ ] <EOL> for oai_id , item in self . _cache [ '<STR_LIT>' ] . items ( ) : <EOL> if oai_id in oai_ids : <EOL> deleted_records . append ( oai_id ) <EOL> item [ '<STR_LIT>' ] = oai_id <EOL> inserted_records . append ( item ) <EOL> for oai_id , item in self . _cache [ '<STR_LIT>' ] . items ( ) : <EOL> if oai_id in oai_ids : <EOL> deleted_sets . append ( oai_id ) <EOL> item [ '<STR_LIT>' ] = oai_id <EOL> inserted_sets . append ( item ) <EOL> for record_id , set_ids in self . _cache [ '<STR_LIT>' ] . items ( ) : <EOL> deleted_setrefs . append ( record_id ) <EOL> for set_id in set_ids : <EOL> inserted_setrefs . append ( <EOL> { '<STR_LIT>' : record_id , '<STR_LIT>' : set_id } ) <EOL> if deleted_records : <EOL> self . _records . delete ( <EOL> self . _records . c . record_id == sql . bindparam ( '<STR_LIT>' ) <EOL> ) . execute ( <EOL> [ { '<STR_LIT>' : rid } for rid in deleted_records ] ) <EOL> if deleted_sets : <EOL> self . _sets . delete ( <EOL> self . _sets . c . set_id == sql . bindparam ( '<STR_LIT>' ) <EOL> ) . execute ( <EOL> [ { '<STR_LIT>' : sid } for sid in deleted_sets ] ) <EOL> if deleted_setrefs : <EOL> self . _setrefs . delete ( <EOL> self . _setrefs . c . record_id == sql . bindparam ( '<STR_LIT>' ) <EOL> ) . execute ( <EOL> [ { '<STR_LIT>' : rid } for rid in deleted_setrefs ] ) <EOL> if inserted_records : <EOL> self . _records . insert ( ) . execute ( inserted_records ) <EOL> if inserted_sets : <EOL> self . _sets . insert ( ) . execute ( inserted_sets ) <EOL> if inserted_setrefs : <EOL> self . _setrefs . insert ( ) . execute ( inserted_setrefs ) <EOL> self . _reset_cache ( ) <EOL> def _reset_cache ( self ) : <EOL> self . _cache = { '<STR_LIT>' : { } , '<STR_LIT>' : { } , '<STR_LIT>' : { } } <EOL> def update_record ( self , oai_id , modified , deleted , sets , metadata ) : <EOL> check_type ( oai_id , <EOL> unicode , <EOL> prefix = "<STR_LIT>" % oai_id , <EOL> suffix = '<STR_LIT>' ) <EOL> check_type ( modified , <EOL> datetime . datetime , <EOL> prefix = "<STR_LIT>" % oai_id , <EOL> suffix = '<STR_LIT>' ) <EOL> check_type ( deleted , <EOL> bool , <EOL> prefix = "<STR_LIT>" % oai_id , <EOL> suffix = '<STR_LIT>' ) <EOL> check_type ( sets , <EOL> dict , <EOL> unicode_values = True , <EOL> recursive = True , <EOL> prefix = "<STR_LIT>" % oai_id , <EOL> suffix = '<STR_LIT>' ) <EOL> check_type ( metadata , <EOL> dict , <EOL> prefix = "<STR_LIT>" % oai_id , <EOL> suffix = '<STR_LIT>' ) <EOL> def date_handler ( obj ) : <EOL> if hasattr ( obj , '<STR_LIT>' ) : <EOL> return obj . isoformat ( ) <EOL> else : <EOL> raise TypeError , '<STR_LIT>' % ( type ( obj ) , repr ( obj ) ) <EOL> metadata = json . dumps ( metadata , default = date_handler ) <EOL> self . _cache [ '<STR_LIT>' ] [ oai_id ] = ( dict ( modified = modified , <EOL> deleted = deleted , <EOL> metadata = metadata ) ) <EOL> self . _cache [ '<STR_LIT>' ] [ oai_id ] = [ ] <EOL> for set_id in sets : <EOL> self . _cache [ '<STR_LIT>' ] [ set_id ] = dict ( <EOL> name = sets [ set_id ] [ '<STR_LIT:name>' ] , <EOL> description = sets [ set_id ] . get ( '<STR_LIT:description>' ) , <EOL> hidden = sets [ set_id ] . get ( '<STR_LIT>' , False ) ) <EOL> self . _cache [ '<STR_LIT>' ] [ oai_id ] . append ( set_id ) <EOL> def get_record ( self , oai_id ) : <EOL> row = self . _records . select ( <EOL> self . _records . c . record_id == oai_id ) . execute ( ) . fetchone ( ) <EOL> if row is None : <EOL> return <EOL> record = { '<STR_LIT:id>' : row . record_id , <EOL> '<STR_LIT>' : row . deleted , <EOL> '<STR_LIT>' : row . modified , <EOL> '<STR_LIT>' : json . loads ( row . metadata ) , <EOL> '<STR_LIT>' : self . get_setrefs ( oai_id ) } <EOL> return record <EOL> def get_set ( self , oai_id ) : <EOL> row = self . _sets . select ( <EOL> self . _sets . c . set_id == oai_id ) . execute ( ) . fetchone ( ) <EOL> if row is None : <EOL> return <EOL> return { '<STR_LIT:id>' : row . set_id , <EOL> '<STR_LIT:name>' : row . name , <EOL> '<STR_LIT:description>' : row . description , <EOL> '<STR_LIT>' : row . hidden } <EOL> def get_setrefs ( self , oai_id , include_hidden_sets = False ) : <EOL> set_ids = [ ] <EOL> query = sql . select ( [ self . _setrefs . c . set_id ] ) <EOL> query . append_whereclause ( self . _setrefs . c . record_id == oai_id ) <EOL> if include_hidden_sets == False : <EOL> query . append_whereclause ( <EOL> sql . and_ ( self . _sets . c . set_id == self . _setrefs . c . set_id , <EOL> self . _sets . c . hidden == include_hidden_sets ) ) <EOL> for row in query . execute ( ) : <EOL> set_ids . append ( row [ <NUM_LIT:0> ] ) <EOL> set_ids . sort ( ) <EOL> return set_ids <EOL> def record_count ( self ) : <EOL> return sql . select ( [ sql . func . count ( '<STR_LIT:*>' ) ] , <EOL> from_obj = [ self . _records ] ) . execute ( ) . fetchone ( ) [ <NUM_LIT:0> ] <EOL> def set_count ( self ) : <EOL> return sql . select ( [ sql . func . count ( '<STR_LIT:*>' ) ] , <EOL> from_obj = [ self . _sets ] ) . execute ( ) . fetchone ( ) [ <NUM_LIT:0> ] <EOL> def remove_record ( self , oai_id ) : <EOL> self . _records . delete ( <EOL> self . _records . c . record_id == oai_id ) . execute ( ) <EOL> self . _setrefs . delete ( <EOL> self . _setrefs . c . record_id == oai_id ) . execute ( ) <EOL> def remove_set ( self , oai_id ) : <EOL> self . _sets . delete ( <EOL> self . _sets . c . set_id == oai_id ) . execute ( ) <EOL> self . _setrefs . delete ( <EOL> self . _setrefs . c . set_id == oai_id ) . execute ( ) <EOL> def oai_sets ( self , offset = <NUM_LIT:0> , batch_size = <NUM_LIT:20> ) : <EOL> for row in self . _sets . select ( <EOL> self . _sets . c . hidden == False <EOL> ) . offset ( offset ) . limit ( batch_size ) . execute ( ) : <EOL> yield { '<STR_LIT:id>' : row . set_id , <EOL> '<STR_LIT:name>' : row . name , <EOL> '<STR_LIT:description>' : row . description } <EOL> def oai_earliest_datestamp ( self ) : <EOL> row = sql . select ( [ self . _records . c . modified ] , <EOL> order_by = [ sql . asc ( self . _records . c . modified ) ] <EOL> ) . limit ( <NUM_LIT:1> ) . execute ( ) . fetchone ( ) <EOL> if row : <EOL> return row [ <NUM_LIT:0> ] <EOL> return datetime . datetime ( <NUM_LIT> , <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> def oai_query ( self , <EOL> offset = <NUM_LIT:0> , <EOL> batch_size = <NUM_LIT:20> , <EOL> needed_sets = None , <EOL> disallowed_sets = None , <EOL> allowed_sets = None , <EOL> from_date = None , <EOL> until_date = None , <EOL> identifier = None ) : <EOL> needed_sets = needed_sets or [ ] <EOL> disallowed_sets = disallowed_sets or [ ] <EOL> allowed_sets = allowed_sets or [ ] <EOL> if batch_size < <NUM_LIT:0> : <EOL> batch_size = <NUM_LIT:0> <EOL> if until_date == None or until_date > datetime . datetime . utcnow ( ) : <EOL> until_date = datetime . datetime . utcnow ( ) <EOL> query = self . _records . select ( <EOL> order_by = [ sql . desc ( self . _records . c . modified ) ] ) <EOL> query . append_whereclause ( self . _records . c . modified <= until_date ) <EOL> if not identifier is None : <EOL> query . append_whereclause ( self . _records . c . record_id == identifier ) <EOL> if not from_date is None : <EOL> query . append_whereclause ( self . _records . c . modified >= from_date ) <EOL> setclauses = [ ] <EOL> for set_id in needed_sets : <EOL> alias = self . _setrefs . alias ( ) <EOL> setclauses . append ( <EOL> sql . and_ ( <EOL> alias . c . set_id == set_id , <EOL> alias . c . record_id == self . _records . c . record_id ) ) <EOL> if setclauses : <EOL> query . append_whereclause ( ( sql . and_ ( * setclauses ) ) ) <EOL> allowed_setclauses = [ ] <EOL> for set_id in allowed_sets : <EOL> alias = self . _setrefs . alias ( ) <EOL> allowed_setclauses . append ( <EOL> sql . and_ ( <EOL> alias . c . set_id == set_id , <EOL> alias . c . record_id == self . _records . c . record_id ) ) <EOL> if allowed_setclauses : <EOL> query . append_whereclause ( sql . or_ ( * allowed_setclauses ) ) <EOL> disallowed_setclauses = [ ] <EOL> for set_id in disallowed_sets : <EOL> alias = self . _setrefs . alias ( ) <EOL> disallowed_setclauses . append ( <EOL> sql . exists ( [ self . _records . c . record_id ] , <EOL> sql . and_ ( <EOL> alias . c . set_id == set_id , <EOL> alias . c . record_id == self . _records . c . record_id ) ) ) <EOL> if disallowed_setclauses : <EOL> query . append_whereclause ( sql . not_ ( sql . or_ ( * disallowed_setclauses ) ) ) <EOL> for row in query . distinct ( ) . offset ( offset ) . limit ( batch_size ) . execute ( ) : <EOL> yield { '<STR_LIT:id>' : row . record_id , <EOL> '<STR_LIT>' : row . deleted , <EOL> '<STR_LIT>' : row . modified , <EOL> '<STR_LIT>' : json . loads ( row . metadata ) , <EOL> '<STR_LIT>' : self . get_setrefs ( row . record_id ) <EOL> } </s>
