rainbowstream/db.py

   1 import os
   2 from sqlalchemy import create_engine
   3 from sqlalchemy.orm import sessionmaker
   4 from sqlalchemy.pool import StaticPool
   5 from .table_def import *
   6
   7
   8 class RainbowDB():
   9
  10     engine = None
  11
  12     def __init__(self):
  13         """
  14         Init DB
  15         """
  16         if os.path.isfile('rainbow.db'):
  17             os.system('rm -rf rainbow.db')
  18         init_db()
  19         self.engine = create_engine('sqlite:///rainbow.db', echo=False, connect_args={'check_same_thread':False}, poolclass=StaticPool)
  20
  21     def tweet_store(self, tweet_id):
  22         """
  23         Store tweet id
  24         """
  25         Session = sessionmaker(bind=self.engine)
  26         session = Session()
  27         t = Tweet(tweet_id)
  28         session.add(t)
  29         session.commit()
  30
  31     def rainbow_to_tweet_query(self, rid):
  32         """
  33         Query base of rainbow id
  34         """
  35         Session = sessionmaker(bind=self.engine)
  36         session = Session()
  37         res = session.query(Tweet).filter_by(rainbow_id=rid).all()
  38         return res
  39
  40     def tweet_to_rainbow_query(self, tid):
  41         """
  42         Query base of tweet id
  43         """
  44         Session = sessionmaker(bind=self.engine)
  45         session = Session()
  46         res = session.query(Tweet).filter_by(tweet_id=tid).all()
  47         return res
  48
  49     def message_store(self, message_id):
  50         """
  51         Store message id
  52         """
  53         Session = sessionmaker(bind=self.engine)
  54         session = Session()
  55         m = Message(message_id)
  56         session.add(m)
  57         session.commit()
  58
  59     def rainbow_to_message_query(self, rid):
  60         """
  61         Query base of rainbow id
  62         """
  63         Session = sessionmaker(bind=self.engine)
  64         session = Session()
  65         res = session.query(Message).filter_by(rainbow_id=rid).all()
  66         return res
  67
  68     def message_to_rainbow_query(self, mid):
  69         """
  70         Query base of message id
  71         """
  72         Session = sessionmaker(bind=self.engine)
  73         session = Session()
  74         res = session.query(Message).filter_by(message_id=mid).all()
  75         return res
  76
  77     def theme_store(self, theme_name):
  78         """
  79         Store theme
  80         """
  81         Session = sessionmaker(bind=self.engine)
  82         session = Session()
  83         th = Theme(theme_name)
  84         session.add(th)
  85         session.commit()
  86
  87     def theme_update(self, theme_name):
  88         """
  89         Update theme
  90         """
  91         Session = sessionmaker(bind=self.engine)
  92         session = Session()
  93         res = session.query(Theme).all()
  94         for r in res:
  95             r.theme_name = theme_name
  96         session.commit()
  97
  98     def theme_query(self):
  99         """
 100         Query theme
 101         """
 102         Session = sessionmaker(bind=self.engine)
 103         session = Session()
 104         res = session.query(Theme).all()
 105         return res
 106
 107     def semaphore_store(self, lock, pause):
 108         """
 109         Store semaphore lock
 110         """
 111         Session = sessionmaker(bind=self.engine)
 112         session = Session()
 113         th = Semaphore(lock, pause)
 114         session.add(th)
 115         session.commit()
 116
 117     def semaphore_update_lock(self, lock):
 118         """
 119         Update semaphore lock
 120         """
 121         Session = sessionmaker(bind=self.engine)
 122         session = Session()
 123         res = session.query(Semaphore).all()
 124         for r in res:
 125             r.lock = lock
 126         session.commit()
 127
 128     def semaphore_update_pause(self, pause):
 129         """
 130         Update semaphore pause
 131         """
 132         Session = sessionmaker(bind=self.engine)
 133         session = Session()
 134         res = session.query(Semaphore).all()
 135         for r in res:
 136             r.pause = pause
 137         session.commit()
 138
 139     def semaphore_query_lock(self):
 140         """
 141         Query semaphore lock
 142         """
 143         Session = sessionmaker(bind=self.engine)
 144         session = Session()
 145         res = session.query(Semaphore).all()
 146         return res[0].lock
 147
 148     def semaphore_query_pause(self):
 149         """
 150         Query semaphore pause
 151         """
 152         Session = sessionmaker(bind=self.engine)
 153         session = Session()
 154         res = session.query(Semaphore).all()
 155         return res[0].pause