refactoring some code
[rainbowstream.git] / 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 def tweet_store(self, tweet_id):
21 """
22 Store tweet id
23 """
24 Session = sessionmaker(bind=self.engine)
25 session = Session()
26 t = Tweet(tweet_id)
27 session.add(t)
28 session.commit()
29
30 def rainbow_to_tweet_query(self, rid):
31 """
32 Query base of rainbow id
33 """
34 Session = sessionmaker(bind=self.engine)
35 session = Session()
36 res = session.query(Tweet).filter_by(rainbow_id=rid).all()
37 return res
38
39 def tweet_to_rainbow_query(self, tid):
40 """
41 Query base of tweet id
42 """
43 Session = sessionmaker(bind=self.engine)
44 session = Session()
45 res = session.query(Tweet).filter_by(tweet_id=tid).all()
46 return res
47
48 def message_store(self, message_id):
49 """
50 Store message id
51 """
52 Session = sessionmaker(bind=self.engine)
53 session = Session()
54 m = Message(message_id)
55 session.add(m)
56 session.commit()
57
58 def rainbow_to_message_query(self, rid):
59 """
60 Query base of rainbow id
61 """
62 Session = sessionmaker(bind=self.engine)
63 session = Session()
64 res = session.query(Message).filter_by(rainbow_id=rid).all()
65 return res
66
67 def message_to_rainbow_query(self, mid):
68 """
69 Query base of message id
70 """
71 Session = sessionmaker(bind=self.engine)
72 session = Session()
73 res = session.query(Message).filter_by(message_id=mid).all()
74 return res
75
76 def theme_store(self, theme_name):
77 """
78 Store theme
79 """
80 Session = sessionmaker(bind=self.engine)
81 session = Session()
82 th = Theme(theme_name)
83 session.add(th)
84 session.commit()
85
86 def theme_update(self, theme_name):
87 """
88 Update theme
89 """
90 Session = sessionmaker(bind=self.engine)
91 session = Session()
92 res = session.query(Theme).all()
93 for r in res:
94 r.theme_name = theme_name
95 session.commit()
96
97 def theme_query(self):
98 """
99 Query theme
100 """
101 Session = sessionmaker(bind=self.engine)
102 session = Session()
103 res = session.query(Theme).all()
104 return res
105
106 def semaphore_store(self, lock, pause):
107 """
108 Store semaphore lock
109 """
110 Session = sessionmaker(bind=self.engine)
111 session = Session()
112 th = Semaphore(lock, pause)
113 session.add(th)
114 session.commit()
115
116 def semaphore_update_lock(self, lock):
117 """
118 Update semaphore lock
119 """
120 Session = sessionmaker(bind=self.engine)
121 session = Session()
122 res = session.query(Semaphore).all()
123 for r in res:
124 r.lock = lock
125 session.commit()
126
127 def semaphore_update_pause(self, pause):
128 """
129 Update semaphore pause
130 """
131 Session = sessionmaker(bind=self.engine)
132 session = Session()
133 res = session.query(Semaphore).all()
134 for r in res:
135 r.pause = pause
136 session.commit()
137
138 def semaphore_query_lock(self):
139 """
140 Query semaphore lock
141 """
142 Session = sessionmaker(bind=self.engine)
143 session = Session()
144 res = session.query(Semaphore).all()
145 return res[0].lock
146
147 def semaphore_query_pause(self):
148 """
149 Query semaphore pause
150 """
151 Session = sessionmaker(bind=self.engine)
152 session = Session()
153 res = session.query(Semaphore).all()
154 return res[0].pause