Setup database

.gitignore

@@ -133,3 +133,6 @@ dmypy.json
 
 # Config file
 config.json
+
+# db
+*.db

bot/__init__.py

@@ -2,8 +2,8 @@ from json import load
 
 from telegram.ext import Updater
 
-
-updater = Updater(token=load(open("config.json"))["bot_token"], use_context=True)
+config = load(open("config.json"))
+updater = Updater(token=config["bot_token"], use_context=True)
 
 import commands
 import callbackQuery

callbackQuery/create.py

@@ -10,7 +10,8 @@ from reportlab.pdfgen import canvas
 from telegram import Update, InlineKeyboardMarkup, InlineKeyboardButton
 from telegram.ext import CallbackContext
 
-from main import reasons, database
+import db
+from main import reasons
 
 local = {
     "work": "travail",
@@ -30,11 +31,14 @@ def create(update: Update, context: CallbackContext):
     reason = reasons[update.effective_chat.id][update["_effective_user"]["id"]]
     del reasons[update.effective_chat.id][update["_effective_user"]["id"]]
     date = datetime.now()
-    first_name = database[update['_effective_user']['id']]['first_name']
-    last_name = database[update['_effective_user']['id']]['last_name']
-    birth_date = database[update['_effective_user']['id']]['birth_date']
-    birth_city = database[update['_effective_user']['id']]['birth_city']
-    address = address_re.fullmatch(database[update['_effective_user']['id']]['address']).groups()
+    s = db.Session()
+    u = s.query(db.User).get(update["_effective_user"]["id"])
+    s.close()
+    first_name = u.first_name
+    last_name = u.last_name
+    birth_date = u.birth_date
+    birth_city = u.birth_city
+    address = address_re.fullmatch(u.address).groups()
 
     img = make(f"Cree le: {date.strftime('%d/%m/%Y a %Hh%M')};\n"
                f"Nom: {first_name};\n"
@@ -54,7 +58,7 @@ def create(update: Update, context: CallbackContext):
     can = canvas.Canvas(packet, pagesize=letter)
     can.setFont("Helvetica", 11)
     can.drawString(119, 696, f"{first_name} {last_name}")
-    can.drawString(119, 674, birth_date)
+    can.drawString(119, 674, birth_date.strftime("%d/%m/%Y"))
     can.drawString(297, 674, birth_city)
     can.drawString(133, 652, f"{address[0]} {address[1]} {address[2]}")
     can.setFontSize(18)

callbackQuery/data.py

@@ -1,22 +1,23 @@
 from telegram import InlineKeyboardButton, InlineKeyboardMarkup, ParseMode, Update
 from telegram.ext import CallbackContext
 
-from main import database
+import db
 
 
 def data(update: Update, context: CallbackContext):
-    if update["_effective_user"]["id"] not in database:
-        database[update["_effective_user"]["id"]] = {"first_name": None,
-                                                     "last_name": None,
-                                                     "birth_date": None,
-                                                     "birth_city": None,
-                                                     "address": None}
+    s = db.Session()
+    u = s.query(db.User).get(update["_effective_user"]["id"])
+    if not u:
+        u = db.User(update["_effective_user"]["id"])
+        s.add(u)
+        s.commit()
+    s.close()
     context.bot.send_message(chat_id=update.effective_chat.id,
-                             text=f"*Firstname*: `{database[update['_effective_user']['id']]['first_name']}`\n"
-                                  f"*Lastname*: `{database[update['_effective_user']['id']]['last_name']}`\n"
-                                  f"*Birth date*: `{database[update['_effective_user']['id']]['birth_date']}`\n"
-                                  f"*Birth city*: `{database[update['_effective_user']['id']]['birth_city']}`\n"
-                                  f"*Address*: `{database[update['_effective_user']['id']]['address']}`\n\n"
+                             text=f"*Firstname*: `{u.first_name}`\n"
+                                  f"*Lastname*: `{u.last_name}`\n"
+                                  f"*Birth date*: `{u.birth_date}`\n"
+                                  f"*Birth city*: `{u.birth_city}`\n"
+                                  f"*Address*: `{u.address}`\n\n"
                                   f"Choose the data you want to edit",
                              parse_mode=ParseMode.MARKDOWN_V2,
                              reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Firstname",

callbackQuery/new.py

@@ -1,11 +1,14 @@
 from telegram import Update, InlineKeyboardMarkup, InlineKeyboardButton, ParseMode
 from telegram.ext import CallbackContext
 
-from main import database, reasons
+import db
+from main import reasons
 
 
 def new(update: Update, context: CallbackContext):
-    if update['_effective_user']['id'] not in database or not all(database[update['_effective_user']['id']].values()):
+    s = db.Session()
+    u = s.query(db.User).get(update["_effective_user"]["id"])
+    if not u or not all([u.first_name, u.last_name, u.birth_date, u.birth_city, u.address]):
         context.bot.send_message(chat_id=update.effective_chat.id, text="You have no data saved !",
                                  reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Set data",
                                                                                           callback_data="data")],

confih_exemple.json

@@ -1,3 +1,4 @@
 {
-  "bot_token": "SUPER_SECRET_TOKEN"
+  "bot_token": "SUPER_SECRET_TOKEN",
+  "db": "postgresql://usr:pass@localhost:5432/sqlalchemy"
 }

db/User.py

@@ -0,0 +1,15 @@
+from db import Base
+from sqlalchemy import Column, Integer, String, Date
+
+
+class User(Base):
+    __tablename__ = "user"
+    id = Column(Integer, primary_key=True)
+    first_name = Column(String)
+    last_name = Column(String)
+    birth_date = Column(Date)
+    birth_city = Column(String)
+    address = Column(String)
+
+    def __init__(self, id: int):
+        self.id = id

db/__init__.py

@@ -0,0 +1,16 @@
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.declarative import declarative_base
+
+from bot import config
+
+
+db = config.get("db")
+args = {"pool_pre_ping": True, "pool_recycle": 3600}
+if not db.startswith("sqlite:"):
+    args.update({"pool_size": 0, "max_overflow": -1})
+engine = create_engine(db, **args)
+Session = sessionmaker(bind=engine)
+Base = declarative_base()
+from db.User import User
+Base.metadata.create_all(engine)

main.py

@@ -8,7 +8,6 @@ local = {"first_name": "Firstname",
          "birth_date": "Birth date",
          "birth_city": "Birth city",
          "address": "Address"}
-database = {}
 messages = {}
 reasons = {}
 

message/edit.py

@@ -4,9 +4,10 @@ from datetime import datetime
 from telegram import Update, ParseMode
 from telegram.ext import CallbackContext
 
+import db
 from callbackQuery.create import address_re
 from callbackQuery.data import data
-from main import local, messages, database
+from main import local, messages
 
 rex = {
     "first_name": re.compile(r"^([a-zA-Z]| )+$"),
@@ -33,6 +34,14 @@ def edit(update: Update, context: CallbackContext, data_edit: str):
         context.bot.send_message(chat_id=update.effective_chat.id, text=f"Invalid value for `{local[name]}` \!",
                                  parse_mode=ParseMode.MARKDOWN_V2)
     else:
+        s = db.Session()
+        u = s.query(db.User).get(update["_effective_user"]["id"])
         del messages[update.effective_chat.id][update["_effective_user"]["id"]]
-        database[update["_effective_user"]["id"]][name] = update.message.text
+        if name == "birth_date":
+            setattr(u, name, datetime.strptime(update.message.text, "%d/%m/%Y").date())
+        else:
+            setattr(u, name, update.message.text)
+        s.add(u)
+        s.commit()
+        s.close()
         data(update, context)

requirements.txt

@@ -12,5 +12,6 @@ pytz==2020.4
 qrcode==6.1
 reportlab==3.5.55
 six==1.15.0
+SQLAlchemy==1.3.20
 tornado==6.1
 tzlocal==2.1