Update README.md
[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 6
b2b933a9 7
18cab06a
O
8class RainbowDB():
9
b2b933a9 10 engine = None
18cab06a
O
11
12 def __init__(self):
eb9781ed
O
13 if not os.path.isfile('rainbow.db'):
14 init_db()
18cab06a
O
15 self.engine = create_engine('sqlite:///rainbow.db', echo=False)
16
17 def store(self, tweet_id):
908d5012 18 """
19 Store tweet id
20 """
18cab06a
O
21 Session = sessionmaker(bind=self.engine)
22 session = Session()
23 m = Map(tweet_id)
24 session.add(m)
25 session.commit()
26
27 def rainbow_query(self, rid):
908d5012 28 """
29 Query base of rainbow id
30 """
18cab06a
O
31 Session = sessionmaker(bind=self.engine)
32 session = Session()
78b81730 33 res = session.query(Map).filter("rainbow_id =:rid").params(rid=rid).all()
18cab06a
O
34 return res
35
36 def tweet_query(self, tid):
908d5012 37 """
38 Query base of tweet id
39 """
18cab06a
O
40 Session = sessionmaker(bind=self.engine)
41 session = Session()
78b81730 42 res = session.query(Map).filter("tweet_id =:tid").params(tid=tid).all()
18cab06a 43 return res
908d5012 44
45 def truncate(self):
46 """
47 Truncate table
48 """
49 Session = sessionmaker(bind=self.engine)
50 session = Session()
51 session.query(Map).delete()
52 session.commit()