bot: test peewee sqlite3

2 files changed, 59 insertions, 34 deletions

diff --git a/dbstore/dbm_store.py b/dbstore/dbm_store.py
deleted file mode 100644
index fedb505..0000000
--- a/dbstore/dbm_store.py
+++ /dev/null
@@ -1,34 +0,0 @@
-import dbm
-
-db = None
-store_file = "fsq_poi.db"
-
-
-def get_loc(fsq_id):
-    global db
-    if db is None:
-        db = dbm.open(store_file, 'c')
-    if fsq_id in db:
-        res = db[fsq_id].decode("utf-8").split("|")
-        return {
-            "name": res[0],
-            "locality": res[1],
-            "region": res[2],
-            "latitude": res[3],
-            "longitude": res[4],
-            "osm_url": res[5],
-        }
-    else:
-        return None
-
-
-def store_loc(loc):
-    global db
-    if db is None:
-        db = dbm.open(store_file, 'c')
-    db[loc["fsq_id"]] = "{}|{}|{}|{}|{}|{}".format(loc["name"],
-                                                   loc["locality"],
-                                                   loc["region"],
-                                                   loc["latitude"],
-                                                   loc["longitude"],
-                                                   loc["osm_url"])
diff --git a/dbstore/peewee_store.py b/dbstore/peewee_store.py
new file mode 100644
index 0000000..37b8c01
--- /dev/null
+++ b/dbstore/peewee_store.py
@@ -0,0 +1,59 @@
+from peewee import *
+
+db = SqliteDatabase("checkinbot.db")
+db.connect(reuse_if_open=True)
+
+
+class BaseModel(Model):
+    class Meta:
+        database = db
+
+
+class User(BaseModel):
+    telegram_user_id = CharField(unique=True, primary_key=True)
+    access_key = CharField(max_length=64)
+    home_instance = CharField(max_length=128)
+    state = CharField(max_length=128)
+    client_id = CharField(max_length=128)
+    client_secret = CharField(max_length=128)
+
+
+class Location(BaseModel):
+    fsq_id = CharField(unique=True, primary_key=True)
+    name = CharField(max_length=128)
+    locality = CharField(max_length=128)
+    region = CharField(max_length=128)
+    latitude = CharField(max_length=128)
+    longitude = CharField(max_length=128)
+
+
+with db.connection_context():
+    db.create_tables([User, Location])
+
+
+def get_poi_by_fsq_id(fsq_id) -> dict:
+    with db.connection_context():
+        try:
+            poi = Location.get(Location.fsq_id == fsq_id)
+            return {
+                "name": poi.name,
+                "locality": poi.locality,
+                "region": poi.region,
+                "latitude": poi.latitude,
+                "longitude": poi.longitude,
+            }
+        except DoesNotExist:
+            return {}
+
+
+def create_poi(poi: dict):
+    with db.connection_context():
+        poi = Location.create(
+            fsq_id=poi["fsq_id"],
+            name=poi["name"],
+            locality=poi["locality"],
+            region=poi["region"],
+            latitude=poi["latitude"],
+            longitude=poi["longitude"],
+        )
+        poi.save()