update
[rainbowstream.git] / rainbowstream / db.py
CommitLineData
eb9781ed 1import os
18cab06a
O
2from sqlalchemy import create_engine
3from sqlalchemy.orm import sessionmaker
4from table_def import Map
eb9781ed 5from .table_def import *
18cab06a
O
6
7class RainbowDB():
8
9 engine=None
10
11 def __init__(self):
eb9781ed
O
12 if not os.path.isfile('rainbow.db'):
13 init_db()
18cab06a
O
14 self.engine = create_engine('sqlite:///rainbow.db', echo=False)
15
16 def store(self, tweet_id):
908d5012 17 """
18 Store tweet id
19 """
18cab06a
O
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):
908d5012 27 """
28 Query base of rainbow id
29 """
18cab06a
O
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):
908d5012 36 """
37 Query base of tweet id
38 """
18cab06a
O
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
908d5012 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()