Source

bubble-economy / persistence.py

Full commit
from sqlalchemy import create_engine, Table, Column, MetaData, BigInteger, SmallInteger, Integer, String, Float, Sequence
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

# http://docs.sqlalchemy.org/en/rel_0_7/core/schema.html#metadata-constraints
# http://docs.sqlalchemy.org/en/rel_0_7/orm/relationships.html
# http://docs.sqlalchemy.org/en/rel_0_7/orm/tutorial.html
# http://docs.sqlalchemy.org/en/rel_0_7/core/tutorial.html
# http://docs.sqlalchemy.org/en/rel_0_7/orm/extensions/declarative.html

Base = declarative_base()

class SummaryStat(Base):
    __tablename__='summary_stat'
    id = Column(Integer, Sequence('summary_stat_id_seq'), primary_key=True, autoincrement=True)
    stat_name = Column(String(256), index=True)
    stat_val = Column(Float)
    noise = Column(Float, index = True)
    dt = Column(Float, index = True)
    radius = Column(Float, index = True)
    base_seed = Column(BigInteger)
    seed = Column(BigInteger)
    discard_steps = Column(Integer)
    n_agents = Column(Integer)
    n_pairs = Column(Integer)
    keep_keys = ['noise', 'dt', 'radius', 'base_seed', 'seed', 'discard_steps', 'n_agents']

def get_db_session(db_url):
    engine = create_engine(db_url)
    Session = sessionmaker(bind=engine)
    session = Session()
    Base.metadata.bind = engine
    Base.metadata.create_all()
    return session

def persist_stat(session, stat_name, stat_val, params=None):
    db_rec = SummaryStat(**dict([(key, params[key]) for key in SummaryStat.keep_keys]))
    summary_stat_rec.stat_name = stat_name
    summary_stat_rec.stat_val = stat_val
    session.add(db_rec)