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         if not os.path.isfile('rainbow.db'):
  13             init_db()
  14         self.engine = create_engine('sqlite:///rainbow.db', echo=False)
  15
  16     def tweet_store(self, tweet_id):
  17         """
  18         Store tweet id
  19         """
  20         Session = sessionmaker(bind=self.engine)
  21         session = Session()
  22         t = Tweet(tweet_id)
  23         session.add(t)
  24         session.commit()
  25
  26     def rainbow_to_tweet_query(self, rid):
  27         """
  28         Query base of rainbow id
  29         """
  30         Session = sessionmaker(bind=self.engine)
  31         session = Session()
  32         res = session.query(Tweet).filter_by(rainbow_id=rid).all()
  33         return res
  34
  35     def tweet_to_rainbow_query(self, tid):
  36         """
  37         Query base of tweet id
  38         """
  39         Session = sessionmaker(bind=self.engine)
  40         session = Session()
  41         res = session.query(Tweet).filter_by(tweet_id=tid).all()
  42         return res
  43
  44     def message_store(self, message_id):
  45         """
  46         Store message id
  47         """
  48         Session = sessionmaker(bind=self.engine)
  49         session = Session()
  50         m = Message(message_id)
  51         session.add(m)
  52         session.commit()
  53
  54     def rainbow_to_message_query(self, rid):
  55         """
  56         Query base of rainbow id
  57         """
  58         Session = sessionmaker(bind=self.engine)
  59         session = Session()
  60         res = session.query(Message).filter_by(rainbow_id=rid).all()
  61         return res
  62
  63     def message_to_rainbow_query(self, mid):
  64         """
  65         Query base of message id
  66         """
  67         Session = sessionmaker(bind=self.engine)
  68         session = Session()
  69         res = session.query(Message).filter_by(message_id=mid).all()
  70         return res
  71
  72     def theme_store(self, theme_name):
  73         """
  74         Store theme
  75         """
  76         Session = sessionmaker(bind=self.engine)
  77         session = Session()
  78         th = Theme(theme_name)
  79         session.add(th)
  80         session.commit()
  81
  82     def theme_update(self, theme_name):
  83         """
  84         Update theme
  85         """
  86         Session = sessionmaker(bind=self.engine)
  87         session = Session()
  88         res = session.query(Theme).all()
  89         for r in res:
  90             r.theme_name = theme_name
  91         session.commit()
  92
  93     def theme_query(self):
  94         """
  95         Query theme
  96         """
  97         Session = sessionmaker(bind=self.engine)
  98         session = Session()
  99         res = session.query(Theme).all()
 100         return res