global by db

[rainbowstream.git] / rainbowstream / db.py

diff --git a/rainbowstream/db.py b/rainbowstream/db.py
new file mode 100644 (file)
index 0000000..d81f5ae
--- /dev/null
+++ b/rainbowstream/db.py
@@ -0,0 +1,30 @@
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from table_def import Map
+
+
+class RainbowDB():
+
+    engine=None
+
+    def __init__(self):
+        self.engine = create_engine('sqlite:///rainbow.db', echo=False)
+
+    def store(self, tweet_id):
+        Session = sessionmaker(bind=self.engine)
+        session = Session()
+        m = Map(tweet_id)
+        session.add(m)
+        session.commit()
+
+    def rainbow_query(self, rid):
+        Session = sessionmaker(bind=self.engine)
+        session = Session()
+        res = session.query(Map).filter("rainbow_id =:rid").params(rid=rid).all()
+        return res
+
+    def tweet_query(self, tid):
+        Session = sessionmaker(bind=self.engine)
+        session = Session()
+        res = session.query(Map).filter("tweet_id =:tid").params(tid=tid).all()
+        return res