Commit | Line | Data |
---|---|---|
eb9781ed | 1 | import os |
18cab06a O |
2 | from sqlalchemy import create_engine |
3 | from sqlalchemy.orm import sessionmaker | |
4 | from table_def import Map | |
eb9781ed | 5 | from .table_def import * |
18cab06a | 6 | |
b2b933a9 | 7 | |
18cab06a O |
8 | class 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() | |
b2b933a9 | 33 | res = session.query(Map) |
34 | .filter("rainbow_id =:rid").params(rid=rid).all() | |
18cab06a O |
35 | return res |
36 | ||
37 | def tweet_query(self, tid): | |
908d5012 | 38 | """ |
39 | Query base of tweet id | |
40 | """ | |
18cab06a O |
41 | Session = sessionmaker(bind=self.engine) |
42 | session = Session() | |
b2b933a9 | 43 | res = session.query(Map) |
44 | .filter("tweet_id =:tid").params(tid=tid).all() | |
18cab06a | 45 | return res |
908d5012 | 46 | |
47 | def truncate(self): | |
48 | """ | |
49 | Truncate table | |
50 | """ | |
51 | Session = sessionmaker(bind=self.engine) | |
52 | session = Session() | |
53 | session.query(Map).delete() | |
54 | session.commit() |