rainbowstream/db.py

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