3108149e |
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 | |
8 | class RainbowDB(): |
9 | |
10 | engine = None |
11 | |
12 | def __init__(self): |
13 | if not os.path.isfile('rainbow.db'): |
14 | init_db() |
15 | self.engine = create_engine('sqlite:///rainbow.db', echo=False) |
16 | |
17 | def store(self, tweet_id): |
18 | """ |
19 | Store tweet id |
20 | """ |
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): |
28 | """ |
29 | Query base of rainbow id |
30 | """ |
31 | Session = sessionmaker(bind=self.engine) |
32 | session = Session() |
33 | res = session.query(Map).filter_by(rainbow_id=rid).all() |
34 | return res |
35 | |
36 | def tweet_query(self, tid): |
37 | """ |
38 | Query base of tweet id |
39 | """ |
40 | Session = sessionmaker(bind=self.engine) |
41 | session = Session() |
42 | res = session.query(Map).filter_by(tweet_id=tid).all() |
43 | return res |