diff --git a/.gitignore b/.gitignore index af3c4fe..46eefd4 100644 --- a/.gitignore +++ b/.gitignore @@ -133,3 +133,6 @@ dmypy.json # Config file config.json + +# db +*.db diff --git a/bot/__init__.py b/bot/__init__.py index 6efa0c6..4adb3ee 100644 --- a/bot/__init__.py +++ b/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 diff --git a/callbackQuery/create.py b/callbackQuery/create.py index 7635fdc..5312a3c 100644 --- a/callbackQuery/create.py +++ b/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) diff --git a/callbackQuery/data.py b/callbackQuery/data.py index 2d50ea0..051fb00 100644 --- a/callbackQuery/data.py +++ b/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", diff --git a/callbackQuery/new.py b/callbackQuery/new.py index 345f76e..e598390 100644 --- a/callbackQuery/new.py +++ b/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")], diff --git a/confih_exemple.json b/confih_exemple.json index 7a0675e..6b5c5e4 100644 --- a/confih_exemple.json +++ b/confih_exemple.json @@ -1,3 +1,4 @@ { - "bot_token": "SUPER_SECRET_TOKEN" + "bot_token": "SUPER_SECRET_TOKEN", + "db": "postgresql://usr:pass@localhost:5432/sqlalchemy" } \ No newline at end of file diff --git a/db/User.py b/db/User.py new file mode 100644 index 0000000..63bf3e3 --- /dev/null +++ b/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 diff --git a/db/__init__.py b/db/__init__.py new file mode 100644 index 0000000..95c53fb --- /dev/null +++ b/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) diff --git a/main.py b/main.py index 6760a7d..beb3988 100644 --- a/main.py +++ b/main.py @@ -8,7 +8,6 @@ local = {"first_name": "Firstname", "birth_date": "Birth date", "birth_city": "Birth city", "address": "Address"} -database = {} messages = {} reasons = {} diff --git a/message/edit.py b/message/edit.py index bdf04be..7ad5637 100644 --- a/message/edit.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index f76097c..135da5b 100644 --- a/requirements.txt +++ b/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