update
[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 Map
5 from .table_def import *
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 store(self, tweet_id):
17 """
18 Store tweet id
19 """
20 Session = sessionmaker(bind=self.engine)
21 session = Session()
22 m = Map(tweet_id)
23 session.add(m)
24 session.commit()
25
26 def rainbow_query(self, rid):
27 """
28 Query base of rainbow id
29 """
30 Session = sessionmaker(bind=self.engine)
31 session = Session()
32 res = session.query(Map).filter("rainbow_id =:rid").params(rid=rid).all()
33 return res
34
35 def tweet_query(self, tid):
36 """
37 Query base of tweet id
38 """
39 Session = sessionmaker(bind=self.engine)
40 session = Session()
41 res = session.query(Map).filter("tweet_id =:tid").params(tid=tid).all()
42 return res
43
44 def truncate(self):
45 """
46 Truncate table
47 """
48 Session = sessionmaker(bind=self.engine)
49 session = Session()
50 session.query(Map).delete()
51 session.commit()