From 3b293151298bed7c1aa32a323779628f894c650b Mon Sep 17 00:00:00 2001 From: kr0sh512 Date: Mon, 23 Dec 2024 04:05:46 +0300 Subject: [PATCH] now on sql --- .gitignore | 3 +- for_db.py | 371 ++++++++++++++++++++++---------------------- postgres_create.sql | 31 ++++ tg.py | 50 ++---- 4 files changed, 232 insertions(+), 223 deletions(-) create mode 100644 postgres_create.sql diff --git a/.gitignore b/.gitignore index d329c12..8a57a7d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ config.yaml database/* __pycache__/* -old.py \ No newline at end of file +old.py +test.py diff --git a/for_db.py b/for_db.py index eb1ef3e..be83281 100644 --- a/for_db.py +++ b/for_db.py @@ -1,161 +1,167 @@ import yaml, os from datetime import datetime, timedelta from telebot import types +from psycopg2 import pool +from sshtunnel import SSHTunnelForwarder + +connection_pool: pool.SimpleConnectionPool = None config = yaml.safe_load(open("config.yaml")) -users_path = config["users_path"] -user_path = config["user_path"] + +server = SSHTunnelForwarder( + (config["host"], 22), + ssh_private_key="~/.ssh/id_rsa", + ssh_username="krosh", + remote_bind_address=("localhost", config["sql_port"]), +) + +server.start() +connection_pool = pool.SimpleConnectionPool( + 1, + 5, + database=config["sql_database"], + user=config["sql_user"], + password=config["sql_password"], + host="localhost", + port=server.local_bind_port, +) + + +def into_dict(curs) -> dict: + cols = [desc[0] for desc in curs.description] + data = curs.fetchone() + + return dict(zip(cols, data)) if data else None + + +def into_list(curs) -> list[dict]: + cols = [desc[0] for desc in curs.description] + data = curs.fetchall() + + return [dict(zip(cols, row)) for row in data] if data else None def check_user(user_id: int) -> dict: - users = yaml.safe_load(open(users_path, "r", encoding="utf-8")) + con = connection_pool.getconn() + curs = con.cursor() + curs.execute("SELECT * FROM users WHERE id = %s", (user_id,)) - if not users: - users = {} + user = into_dict(curs) + connection_pool.putconn(con) - if user_id in users: - return users[user_id] - - return None + return user def new_user(message: types.Message) -> bool: if check_user(message.chat.id): return False - update_user_settings(message.chat.id, "id", message.chat.id) - update_user_settings(message.chat.id, "type", message.chat.type) - update_user_settings(message.chat.id, "username", message.chat.username) - update_user_settings(message.chat.id, "first_name", message.chat.first_name) - update_user_settings(message.chat.id, "last_name", message.chat.last_name) - update_user_settings( - message.chat.id, "time_created", datetime.now().strftime("%Y-%m-%d %H:%M:%S") + con = connection_pool.getconn() + curs = con.cursor() + curs.execute( + """ + INSERT INTO users (id, username, firstname, lastname, chat) + VALUES (%s, %s, %s, %s, %s) + """, + ( + message.chat.id, + message.chat.username, + message.chat.first_name, + message.chat.last_name, + message.chat.type, + ), ) - update_user_settings(message.chat.id, "last_message", None) - update_user_settings( - message.chat.id, "remind_delta", 12 * 60 * 60 - ) # 12 часов по умолчанию + con.commit() + connection_pool.putconn(con) return def add_note( - id: int, text: str, time_notif: datetime = None, list: str = "Default" + id: int, text: str, remind_at: datetime = None, list: str = "Default" ) -> bool: if not check_user(id): return False - notes = ( - yaml.safe_load(open(user_path.format(id), "r", encoding="utf-8")) - if os.path.exists(user_path.format(id)) - else {} + con = connection_pool.getconn() + curs = con.cursor() + curs.execute( + """ + SELECT id FROM lists WHERE name = %s AND user_id = %s + """, + (list, id), ) - if list not in notes: - notes[list] = {"name": list, "description": "Стандартный список", "notes": []} + data = curs.fetchone() - notes[list]["notes"].append( - { - "text": text, - "time_notif": str(time_notif) if time_notif else None, - "time_created": str(datetime.now().strftime("%Y-%m-%d %H:%M:%S")), - } - ) - - if True: # TODO: edit sorting - import functools - - def comp(x, y): - if not x["time_notif"] and not y["time_notif"]: - return 0 - elif not x["time_notif"]: - return -1 - elif not y["time_notif"]: - return 1 - else: - x_time = datetime.strptime( - x["time_notif"], - "%Y-%m-%d %H:%M:%S", - ) - y_time = datetime.strptime( - y["time_notif"], - "%Y-%m-%d %H:%M:%S", - ) - if x_time > y_time: - return 1 - elif x_time < y_time: - return -1 - return 0 - - notes[list]["notes"].sort( - key=lambda x: datetime.strptime( - x["time_created"], - "%Y-%m-%d %H:%M:%S", - ), - reverse=True, + if not data: + curs.execute( + """ + INSERT INTO lists (name, user_id) + VALUES (%s, %s) + """, + (list, id), ) + con.commit() - notes[list]["notes"].sort( - key=functools.cmp_to_key(comp), - ) - - yaml_data = yaml.dump( - notes, - default_flow_style=False, - encoding="utf-8", - allow_unicode=True, - width=float("inf"), - sort_keys=False, + curs.execute( + """ + INSERT INTO notes (user_id, list_id, content, remind_at) + VALUES (%s, (SELECT id FROM lists WHERE name = %s AND user_id = %s), %s, %s) + """, + (id, list, id, text, remind_at), ) + con.commit() + rowcnt = curs.rowcount - with open(user_path.format(id), "wb") as file: - file.write(yaml_data) + connection_pool.putconn(con) - return True + return True if rowcnt else False -def get_notes(id: int, list="Default") -> dict: +def get_notes(id: int, list="Default") -> list[dict]: if not check_user(id): return None - if not os.path.exists(user_path.format(id)): - return None + con = connection_pool.getconn() + curs = con.cursor() - notes = yaml.safe_load(open(user_path.format(id), "r", encoding="utf-8")) + curs.execute( + """ + SELECT * FROM notes + WHERE user_id = %s AND list_id = (SELECT id FROM lists WHERE name = %s AND user_id = %s) + ORDER BY remind_at IS NULL DESC, remind_at ASC, created_at DESC + """, + (id, list, id), + ) - return notes[list] if list in notes else None + notes = into_list(curs) + connection_pool.putconn(con) + + return notes def delete_note(id, note_ind, list="Default") -> bool: if not check_user(id): return False - if not os.path.exists(user_path.format(id)): - return False - - notes = yaml.safe_load(open(user_path.format(id), "r", encoding="utf-8")) - - if list not in notes: - return False - - if note_ind >= len(notes[list]["notes"]): - return False - - notes[list]["notes"].pop(note_ind) - - yaml_data = yaml.dump( - notes, - default_flow_style=False, - encoding="utf-8", - allow_unicode=True, - width=float("inf"), - sort_keys=False, + con = connection_pool.getconn() + curs = con.cursor() + curs.execute( + """ + DELETE FROM notes + WHERE user_id = %s AND list_id = (SELECT id FROM lists WHERE name = %s AND user_id = %s) + ORDER BY remind_at IS NULL DESC, remind_at ASC, created_at DESC + LIMIT 1 OFFSET %s + """, + (id, list, id, note_ind), ) - with open(user_path.format(id), "wb") as file: - file.write(yaml_data) + con.commit() + rowcnt = curs.rowcount + connection_pool.putconn(con) - return True + return True if rowcnt else False def new_message(message: types.Message) -> int: @@ -165,105 +171,98 @@ def new_message(message: types.Message) -> int: if not check_user(id): return None - users = yaml.safe_load(open(users_path, "r", encoding="utf-8")) - - user = users[id] - - if "last_message" not in user: - user["last_message"] = None - - last_id = user["last_message"] - user["last_message"] = new_msg_id - - users[id] = user - - yaml_data = yaml.dump( - users, - default_flow_style=False, - encoding="utf-8", - allow_unicode=True, - width=float("inf"), - sort_keys=False, + con = connection_pool.getconn() + curs = con.cursor() + curs.execute( + """ + SELECT last_message FROM users WHERE id = %s + """, + (id,), ) + last_id = curs.fetchone()[0] - with open(users_path, "wb") as file: - file.write(yaml_data) + curs.execute( + """ + UPDATE users + SET last_message = %s + WHERE id = %s + """, + (new_msg_id, id), + ) + con.commit() + connection_pool.putconn(con) return last_id def check_old_notes() -> list[int, str, int]: # возвращает первую устаревшую заметку - users = yaml.safe_load(open(users_path, "r", encoding="utf-8")) + con = connection_pool.getconn() + curs = con.cursor() + curs.execute( + """ + SELECT id FROM users + """ + ) + users = curs.fetchall() - if not users: - users = {} + for user_id in users: + curs.execute( + """ + SELECT * FROM notes + WHERE id = %s + ORDER BY remind_at IS NULL DESC, remind_at ASC, created_at DESC + """, + (user_id,), + ) + notes = [dict(row) for row in curs.fetchall()] - for user_id in users.keys(): - if not os.path.exists(user_path.format(user_id)): + if not notes: continue - notes = yaml.safe_load(open(user_path.format(user_id), "r", encoding="utf-8")) - if "remind_delta" not in users[user_id]: - users[user_id]["remind_delta"] = 12 * 60 * 60 # не сохранится! + timedt = timedelta(seconds=user_settings(user_id)["remind_delta"]) - timedt = timedelta(seconds=users[user_id]["remind_delta"]) + for ind in range(len(notes)): + if ( + notes[ind]["remind_at"] + and datetime.strptime(notes[ind]["remind_at"], "%Y-%m-%d %H:%M:%S") + - timedt + < datetime.now() + ): + curs.execute( + """ + SELECT name FROM lists + WHERE id = %s + """, + (notes[ind]["list_id"],), + ) + list_name = curs.fetchone() + connection_pool.putconn(con) - for list in notes.keys(): - for ind in range(len(notes[list]["notes"])): - if ( - notes[list]["notes"][ind]["time_notif"] - and datetime.strptime( - notes[list]["notes"][ind]["time_notif"], "%Y-%m-%d %H:%M:%S" - ) - - timedt - < datetime.now() - ): - return user_id, list, ind + return (user_id, list_name, ind) + + connection_pool.putconn(con) return None, None, None def update_user_settings(user_id: int, param: str, value: any) -> bool: - users = yaml.safe_load(open(users_path, "r", encoding="utf-8")) + con = connection_pool.getconn() + curs = con.cursor() + query = f"UPDATE users SET {param} = %s WHERE id = %s" + curs.execute(query, (value, user_id)) + con.commit() + rowcnt = curs.rowcount + connection_pool.putconn(con) - if not users: - users = {} - - if user_id not in users: - users[user_id] = {} - - users[user_id][param] = value - - yaml_data = yaml.dump( - users, - default_flow_style=False, - encoding="utf-8", - allow_unicode=True, - width=float("inf"), - sort_keys=False, - ) - - with open(users_path, "wb") as file: - file.write(yaml_data) - - with open(user_path.format(user_id), "wb") as file: - yaml_data = yaml.dump( - {}, - default_flow_style=False, - encoding="utf-8", - allow_unicode=True, - width=float("inf"), - sort_keys=False, - ) - file.write(yaml_data) - - return True + return True if rowcnt else False def user_settings(user_id: int) -> dict: - if not check_user(user_id): - return None + con = connection_pool.getconn() + curs = con.cursor() + curs.execute("SELECT * FROM users WHERE id = %s", (user_id,)) - users = yaml.safe_load(open(users_path, "r", encoding="utf-8")) + user = into_dict(curs) + connection_pool.putconn(con) - return users[user_id] + return user diff --git a/postgres_create.sql b/postgres_create.sql new file mode 100644 index 0000000..990e607 --- /dev/null +++ b/postgres_create.sql @@ -0,0 +1,31 @@ +CREATE type chat_type AS ENUM ('private', 'group', 'supergroup', 'channel'); + +CREATE TABLE users ( + id BIGINT PRIMARY KEY, + username TEXT, + firstname TEXT, + lastname TEXT, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + last_message BIGINT NOT NULL DEFAULT 0, + remind_delta BIGINT NOT NULL DEFAULT 21600, + chat chat_type NOT NULL +); + +CREATE Table lists ( + id BIGSERIAL PRIMARY KEY, + user_id BIGINT NOT NULL REFERENCES users (id), + name TEXT NOT NULL, + discription TEXT, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP +); + + +CREATE TABLE notes ( + id BIGSERIAL PRIMARY KEY, + user_id BIGINT NOT NULL REFERENCES users (id), + list_id BIGINT NOT NULL REFERENCES lists (id), + content TEXT NOT NULL, + remind_at TIMESTAMP, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP +); + diff --git a/tg.py b/tg.py index ae19c7f..c8231b9 100644 --- a/tg.py +++ b/tg.py @@ -11,7 +11,7 @@ config = yaml.safe_load(open("config.yaml")) lang = yaml.safe_load(open("lang.yaml", encoding="utf-8"))["ru"] list_reactions = yaml.safe_load(open("lang.yaml", encoding="utf-8"))["list_reaction"] bot = telebot.TeleBot( - config["api_token"], + config["test_token"], colorful_logs=True, disable_web_page_preview=True, parse_mode=ParseMode.HTML, @@ -20,7 +20,6 @@ bot = telebot.TeleBot( admin_id = config["admin_id"] bot_username = config["bot_username"] - locale.setlocale(locale.LC_ALL, "ru_RU.UTF-8") @@ -77,13 +76,13 @@ def delete_notes_callback(call): note_ind = int(call.data.split("#")[2]) prefix = call.data.split("#")[0] - note = db.get_notes(call.message.chat.id, prefix)["notes"][note_ind] + note = db.get_notes(call.message.chat.id, prefix)[note_ind] db.delete_note(call.message.chat.id, note_ind, prefix) list_notes(call.message, prefix, edit=True) - bot.send_message(call.message.chat.id, f"❌ Удалено:\n\n{note['text']}") + bot.send_message(call.message.chat.id, f"❌ Удалено:\n\n{note['content']}") return @@ -92,7 +91,7 @@ def delete_notes_callback(call): def edit_notes_callback(call): note_ind = int(call.data.split("#")[2]) prefix = call.data.split("#")[0] - notes = db.get_notes(call.message.chat.id, prefix)["notes"] + notes = db.get_notes(call.message.chat.id, prefix) note = notes[note_ind] @@ -112,10 +111,10 @@ def edit_notes_callback(call): text = f"📝 Заметка {note_ind + 1}" - if note["time_notif"]: - text += f" ⌚️ {datetime.strptime(note['time_notif'], '%Y-%m-%d %H:%M:%S').strftime('%d %B %Y')}" + if note["remind_at"]: + text += f" ⌚️ {note['remind_at'].strftime('%d %B %Y')}" - text += f"\n\n{note['text']}" + text += f"\n\n{note['content']}" bot.edit_message_text( text, @@ -133,7 +132,7 @@ def choose_note_callback(call): prefix = call.data.split("#")[0] markup = types.InlineKeyboardMarkup() num_on_page = 6 # Лимит телеграмма - 8 - notes = db.get_notes(call.message.chat.id, prefix)["notes"] + notes = db.get_notes(call.message.chat.id, prefix) max_page = (len(notes) - 1) // num_on_page lst_button = [] @@ -256,25 +255,6 @@ def display_settings(message: types.Message): return - settings = db.user_settings(message.chat.id) - # if not settings: - # send_message(message, lang["not_registered"]) - - # return - - settings_msg = "⚙️ Настройки:\ - \n\nВремя до отправки уведомления: {}" - - markup = types.InlineKeyboardMarkup() - markup.add( - types.InlineKeyboardButton( - "Изменить время", - callback_data="edit_time#", - ) - ) - - return - @bot.message_handler(commands=["list", "lists", "l"]) def list_notes(message: types.Message, list: str = "Default", edit: bool = False): @@ -285,15 +265,13 @@ def list_notes(message: types.Message, list: str = "Default", edit: bool = False return - list_notes = notes["notes"] - notes_msg = "⚡️ Ваши напоминания:\n\n" - for i in range(len(list_notes)): + for i in range(len(notes)): notes_msg += f"{i + 1}) " - if list_notes[i]["time_notif"]: - notes_msg += f" ⌚️ {datetime.strptime(list_notes[i]['time_notif'], '%Y-%m-%d %H:%M:%S').strftime('%d %B %Y')}:\n" + if notes[i]["remind_at"]: + notes_msg += f" ⌚️ {notes[i]['remind_at'].strftime('%d %B %Y')}:\n" - notes_msg += f"{list_notes[i]['text']}\n\n" + notes_msg += f"{notes[i]['content']}\n\n" markup = types.InlineKeyboardMarkup() markup.add( @@ -467,13 +445,13 @@ if __name__ == "__main__": chat_id, list, ind_note = db.check_old_notes() if chat_id: - note = db.get_notes(chat_id, list)["notes"][ind_note] + note = db.get_notes(chat_id, list)[ind_note] db.delete_note(chat_id, ind_note, list) msg = bot.send_message( chat_id, - f"❌ {note['text']} \n(Уведомление: {datetime.strptime(note['time_notif'], '%Y-%m-%d %H:%M:%S').strftime('%d %B %Y')})", + f"❌ {note['content']} \n(Уведомление: {note['remind_at'].strftime('%d %B %Y')})", timeout=1, )