<s> from contextlib import contextmanager <EOL> import logging <EOL> import sqlalchemy <EOL> from sqlalchemy import and_ <EOL> from sqlalchemy . exc import DisconnectionError <EOL> from sqlalchemy . exc import IntegrityError <EOL> from sqlalchemy . exc import OperationalError <EOL> from sqlalchemy . orm . exc import NoResultFound <EOL> from sqlalchemy . orm import sessionmaker <EOL> from sqlalchemy . sql import select <EOL> from winchester . config import ConfigItem <EOL> from winchester . config import ConfigManager <EOL> from winchester import models <EOL> logger = logging . getLogger ( __name__ ) <EOL> ENGINES = dict ( ) <EOL> SESSIONMAKERS = dict ( ) <EOL> class DatabaseConnectionError ( models . DBException ) : <EOL> pass <EOL> class DuplicateError ( models . DBException ) : <EOL> pass <EOL> class LockError ( models . DBException ) : <EOL> pass <EOL> class NoSuchEventError ( models . DBException ) : <EOL> pass <EOL> class NoSuchStreamError ( models . DBException ) : <EOL> pass <EOL> def sessioned ( func ) : <EOL> def with_session ( self , * args , ** kw ) : <EOL> if '<STR_LIT>' in kw : <EOL> return func ( self , * args , ** kw ) <EOL> else : <EOL> with self . in_session ( ) as session : <EOL> kw [ '<STR_LIT>' ] = session <EOL> retval = func ( self , * args , ** kw ) <EOL> return retval <EOL> return with_session <EOL> class DBInterface ( object ) : <EOL> @ classmethod <EOL> def config_description ( cls ) : <EOL> return dict ( <EOL> url = ConfigItem ( required = True , <EOL> help = "<STR_LIT>" ) , <EOL> ) <EOL> def __init__ ( self , config ) : <EOL> self . config = ConfigManager . wrap ( config , self . config_description ( ) ) <EOL> self . db_url = config [ '<STR_LIT:url>' ] <EOL> self . echo_sql = config . get ( '<STR_LIT>' , False ) <EOL> @ property <EOL> def engine ( self ) : <EOL> global ENGINES <EOL> if self . db_url not in ENGINES : <EOL> engine = sqlalchemy . create_engine ( self . db_url , echo = self . echo_sql ) <EOL> ENGINES [ self . db_url ] = engine <EOL> return ENGINES [ self . db_url ] <EOL> @ property <EOL> def sessionmaker ( self ) : <EOL> global SESSIONMAKERS <EOL> if self . db_url not in SESSIONMAKERS : <EOL> maker = sessionmaker ( bind = self . engine ) <EOL> SESSIONMAKERS [ self . db_url ] = maker <EOL> return SESSIONMAKERS [ self . db_url ] <EOL> def close ( self ) : <EOL> if self . db_url in ENGINES : <EOL> del ENGINES [ self . db_url ] <EOL> if self . db_url in SESSIONMAKERS : <EOL> del SESSIONMAKERS [ self . db_url ] <EOL> def get_session ( self ) : <EOL> return self . sessionmaker ( expire_on_commit = False ) <EOL> @ contextmanager <EOL> def in_session ( self ) : <EOL> """<STR_LIT>""" <EOL> session = self . get_session ( ) <EOL> try : <EOL> yield session <EOL> session . commit ( ) <EOL> except IntegrityError : <EOL> session . rollback ( ) <EOL> raise DuplicateError ( "<STR_LIT>" ) <EOL> except ( OperationalError , DisconnectionError ) : <EOL> session . rollback ( ) <EOL> self . close ( ) <EOL> logger . warn ( "<STR_LIT>" ) <EOL> raise DatabaseConnectionError ( ) <EOL> except Exception : <EOL> session . rollback ( ) <EOL> raise <EOL> finally : <EOL> session . close ( ) <EOL> @ sessioned <EOL> def get_event_type ( self , description , session = None ) : <EOL> t = session . query ( models . EventType ) . filter ( <EOL> models . EventType . desc == description ) . first ( ) <EOL> if t is None : <EOL> t = models . EventType ( description ) <EOL> session . add ( t ) <EOL> return t <EOL> @ sessioned <EOL> def create_event ( self , message_id , event_type , generated , traits , <EOL> session = None ) : <EOL> event_type = self . get_event_type ( event_type , session = session ) <EOL> e = models . Event ( message_id , event_type , generated ) <EOL> for name in traits : <EOL> try : <EOL> e [ name ] = traits [ name ] <EOL> except models . InvalidTraitType : <EOL> logger . error ( "<STR_LIT>" <EOL> "<STR_LIT>" % ( name , traits [ name ] , <EOL> type ( traits [ name ] ) ) ) <EOL> session . add ( e ) <EOL> @ sessioned <EOL> def get_event_by_message_id ( self , message_id , session = None ) : <EOL> try : <EOL> e = session . query ( models . Event ) . filter ( models . Event . message_id == message_id ) . one ( ) <EOL> except NoResultFound : <EOL> raise NoSuchEventError ( <EOL> "<STR_LIT>" % message_id ) <EOL> return e . as_dict <EOL> @ sessioned <EOL> def find_older_events ( self , purge_date , batchsize , session = None ) : <EOL> conn = session . connection ( ) <EOL> event_table = models . Event . __table__ <EOL> q = select ( [ event_table . c . id ] ) <EOL> q = q . where ( event_table . c . generated < purge_date ) <EOL> q = q . order_by ( event_table . c . generated . asc ( ) ) <EOL> q = q . limit ( batchsize ) <EOL> return [ r [ <NUM_LIT:0> ] for r in conn . execute ( q ) . fetchall ( ) ] <EOL> @ sessioned <EOL> def purge_events ( self , event_ids , session = None ) : <EOL> conn = session . connection ( ) <EOL> dq = models . stream_event_table . delete ( ) <EOL> dq = dq . where ( models . stream_event_table . c . event_id . in_ ( event_ids ) ) <EOL> conn . execute ( dq ) <EOL> trait_table = models . Trait . __table__ <EOL> dq = trait_table . delete ( ) <EOL> dq = dq . where ( trait_table . c . event_id . in_ ( event_ids ) ) <EOL> conn . execute ( dq ) <EOL> event_table = models . Event . __table__ <EOL> dq = event_table . delete ( ) <EOL> dq = dq . where ( event_table . c . id . in_ ( event_ids ) ) <EOL> conn . execute ( dq ) <EOL> @ sessioned <EOL> def find_events ( self , from_datetime = None , to_datetime = None , <EOL> event_name = None , traits = None , mark = None , limit = None , <EOL> session = None , count = False ) : <EOL> order_desc = True <EOL> q = session . query ( models . Event ) <EOL> if mark is not None : <EOL> if mark . startswith ( '<STR_LIT:+>' ) : <EOL> order_desc = False <EOL> mark = mark [ <NUM_LIT:1> : ] <EOL> if mark . startswith ( '<STR_LIT:->' ) : <EOL> order_desc = True <EOL> mark = mark [ <NUM_LIT:1> : ] <EOL> if mark : <EOL> if order_desc : <EOL> q = q . filter ( models . Event . id < int ( mark , <NUM_LIT:16> ) ) <EOL> else : <EOL> q = q . filter ( models . Event . id > int ( mark , <NUM_LIT:16> ) ) <EOL> if from_datetime is not None : <EOL> q = q . filter ( models . Event . generated > from_datetime ) <EOL> if to_datetime is not None : <EOL> q = q . filter ( models . Event . generated <= to_datetime ) <EOL> if event_name is not None : <EOL> event_type = self . get_event_type ( event_name , <EOL> session = session ) <EOL> q = q . filter ( models . Event . event_type_id == event_type . id ) <EOL> if traits is not None : <EOL> for name , val in traits . items ( ) : <EOL> q = q . filter ( models . Event . traits . any ( and_ ( <EOL> models . Trait . name == name , <EOL> models . Trait . value == val ) ) ) <EOL> if count : <EOL> q = q . count ( ) <EOL> return [ { "<STR_LIT:count>" : q } ] <EOL> if order_desc : <EOL> q = q . order_by ( models . Event . id . desc ( ) ) <EOL> mark_fmt = '<STR_LIT>' <EOL> else : <EOL> q = q . order_by ( models . Event . id . asc ( ) ) <EOL> mark_fmt = '<STR_LIT>' <EOL> if limit is not None : <EOL> q = q . limit ( limit ) <EOL> event_info = [ ] <EOL> for event in q . all ( ) : <EOL> info = event . as_dict <EOL> info [ '<STR_LIT>' ] = mark_fmt % event . id <EOL> event_info . append ( info ) <EOL> return event_info <EOL> @ sessioned <EOL> def get_stream_by_id ( self , stream_id , session = None ) : <EOL> try : <EOL> s = session . query ( models . Stream ) . filter ( models . Stream . id == stream_id ) . one ( ) <EOL> except NoResultFound : <EOL> raise NoSuchStreamError ( "<STR_LIT>" % stream_id ) <EOL> return s <EOL> @ sessioned <EOL> def create_stream ( self , trigger_name , initial_event , dist_traits , <EOL> expire_expr , session = None ) : <EOL> first_event_time = initial_event [ '<STR_LIT>' ] <EOL> s = models . Stream ( trigger_name , first_event_time ) <EOL> for trait_name in dist_traits : <EOL> s [ trait_name ] = dist_traits [ trait_name ] <EOL> session . add ( s ) <EOL> self . add_event_stream ( s , initial_event , expire_expr , session = session ) <EOL> return s <EOL> @ sessioned <EOL> def stream_has_dist_trait ( self , stream_id , name , value = None , session = None ) : <EOL> q = session . query ( models . DistinguishingTrait ) <EOL> q = q . filter ( models . DistinguishingTrait . stream_id == stream_id ) <EOL> q = q . filter ( models . DistinguishingTrait . name == name ) <EOL> if value is not None : <EOL> q = q . filter ( models . DistinguishingTrait . value == value ) <EOL> dt = q . first ( ) <EOL> if dt is not None : <EOL> dt = dt . as_dict <EOL> return dt <EOL> @ sessioned <EOL> def get_stream_events ( self , stream , session = None ) : <EOL> if stream not in session : <EOL> stream = session . merge ( stream ) <EOL> return [ event . as_dict for event in stream . events ] <EOL> @ sessioned <EOL> def add_event_stream ( self , stream , event , expire_expr , session = None ) : <EOL> if stream not in session : <EOL> session . add ( stream ) <EOL> message_id = event [ '<STR_LIT>' ] <EOL> timestamp = event [ '<STR_LIT>' ] <EOL> if timestamp < stream . first_event : <EOL> stream . first_event = timestamp <EOL> if timestamp > stream . last_event : <EOL> stream . last_event = timestamp <EOL> stream . expire_timestamp = expire_expr ( first = stream . first_event , <EOL> last = stream . last_event ) . timestamp <EOL> eq = session . query ( models . Event ) <EOL> eq = eq . filter ( models . Event . message_id == message_id ) <EOL> e = eq . one ( ) <EOL> stream . events . append ( e ) <EOL> return e <EOL> @ sessioned <EOL> def get_active_stream ( self , name , dist_traits , current_time , session = None ) : <EOL> q = session . query ( models . Stream ) <EOL> q = q . filter ( models . Stream . name == name ) <EOL> q = q . filter ( models . Stream . state == int ( models . StreamState . active ) ) <EOL> q = q . filter ( models . Stream . expire_timestamp > current_time ) <EOL> for name , val in dist_traits . items ( ) : <EOL> q = q . filter ( models . Stream . distinguished_by . any ( and_ ( <EOL> models . DistinguishingTrait . name == name , <EOL> models . DistinguishingTrait . value == val ) ) ) <EOL> return q . first ( ) <EOL> @ sessioned <EOL> def stream_ready_to_fire ( self , stream , timestamp , session = None ) : <EOL> if stream not in session : <EOL> session . add ( stream ) <EOL> stream . fire_timestamp = timestamp <EOL> @ sessioned <EOL> def get_ready_streams ( self , batch_size , current_time , expire = False , <EOL> session = None ) : <EOL> q = session . query ( models . Stream ) <EOL> if expire : <EOL> states = ( int ( models . StreamState . active ) , <EOL> int ( models . StreamState . retry_expire ) ) <EOL> else : <EOL> states = ( <EOL> int ( models . StreamState . active ) , <EOL> int ( models . StreamState . retry_fire ) ) <EOL> q = q . filter ( models . Stream . state . in_ ( states ) ) <EOL> if expire : <EOL> q = q . filter ( models . Stream . expire_timestamp < current_time ) <EOL> else : <EOL> q = q . filter ( models . Stream . fire_timestamp < current_time ) <EOL> q = q . limit ( batch_size ) <EOL> return q . all ( ) <EOL> def set_stream_state ( self , stream , state ) : <EOL> serial = stream . state_serial_no <EOL> stream_id = stream . id <EOL> with self . in_session ( ) as session : <EOL> q = session . query ( models . Stream ) <EOL> q = q . filter ( models . Stream . id == stream_id ) <EOL> q = q . filter ( models . Stream . state_serial_no == serial ) <EOL> ct = q . update ( dict ( state = int ( state ) , state_serial_no = serial + <NUM_LIT:1> ) ) <EOL> if ct != <NUM_LIT:1> : <EOL> raise LockError ( "<STR_LIT>" ) <EOL> return self . get_stream_by_id ( stream_id ) <EOL> def reset_stream ( self , stream ) : <EOL> if stream . state == models . StreamState . error : <EOL> return self . set_stream_state ( stream , models . StreamState . retry_fire ) <EOL> if stream . state == models . StreamState . expire_error : <EOL> return self . set_stream_state ( stream , <EOL> models . StreamState . retry_expire ) <EOL> return stream <EOL> @ sessioned <EOL> def find_streams ( self , count = False , stream_id = None , state = None , <EOL> older_than = None , younger_than = None , <EOL> name = None , distinguishing_traits = None , <EOL> session = None , include_events = False , <EOL> limit = None , mark = None ) : <EOL> order_desc = True <EOL> q = session . query ( models . Stream ) <EOL> if mark is not None : <EOL> if mark . startswith ( '<STR_LIT:+>' ) : <EOL> order_desc = False <EOL> mark = mark [ <NUM_LIT:1> : ] <EOL> if mark . startswith ( '<STR_LIT:->' ) : <EOL> order_desc = True <EOL> mark = mark [ <NUM_LIT:1> : ] <EOL> if mark : <EOL> if order_desc : <EOL> q = q . filter ( models . Stream . id < int ( mark , <NUM_LIT:16> ) ) <EOL> else : <EOL> q = q . filter ( models . Stream . id > int ( mark , <NUM_LIT:16> ) ) <EOL> if stream_id is not None : <EOL> q = q . filter ( models . Stream . id == stream_id ) <EOL> if state is not None : <EOL> q = q . filter ( models . Stream . state == int ( state ) ) <EOL> if older_than is not None : <EOL> q = q . filter ( models . Stream . first_event < older_than ) <EOL> if younger_than is not None : <EOL> q = q . filter ( models . Stream . last_event > younger_than ) <EOL> if name is not None : <EOL> q = q . filter ( models . Stream . name == name ) <EOL> if distinguishing_traits is not None : <EOL> for name , val in distinguishing_traits . items ( ) : <EOL> q = q . filter ( models . Stream . distinguished_by . any ( and_ ( <EOL> models . DistinguishingTrait . name == name , <EOL> models . DistinguishingTrait . value == val ) ) ) <EOL> if count : <EOL> q = q . count ( ) <EOL> return [ { "<STR_LIT:count>" : q } ] <EOL> if order_desc : <EOL> q = q . order_by ( models . Stream . id . desc ( ) ) <EOL> mark_fmt = '<STR_LIT>' <EOL> else : <EOL> q = q . order_by ( models . Stream . id . asc ( ) ) <EOL> mark_fmt = '<STR_LIT>' <EOL> if limit is not None : <EOL> q = q . limit ( limit ) <EOL> stream_info = [ ] <EOL> for stream in q . all ( ) : <EOL> info = stream . as_dict <EOL> info [ '<STR_LIT>' ] = mark_fmt % stream . id <EOL> if include_events : <EOL> info [ '<STR_LIT>' ] = self . get_stream_events ( <EOL> stream , session = session ) <EOL> stream_info . append ( info ) <EOL> return stream_info <EOL> @ sessioned <EOL> def purge_stream ( self , stream , session = None ) : <EOL> if stream not in session : <EOL> session . add ( stream ) <EOL> session . delete ( stream ) </s>
