now on sql

This commit is contained in:
kr0sh512
2024-12-23 04:05:46 +03:00
parent 0e7c185517
commit 3b29315129
4 changed files with 232 additions and 223 deletions
+2 -1
View File
@@ -1,4 +1,5 @@
config.yaml config.yaml
database/* database/*
__pycache__/* __pycache__/*
old.py old.py
test.py
+185 -186
View File
@@ -1,161 +1,167 @@
import yaml, os import yaml, os
from datetime import datetime, timedelta from datetime import datetime, timedelta
from telebot import types from telebot import types
from psycopg2 import pool
from sshtunnel import SSHTunnelForwarder
connection_pool: pool.SimpleConnectionPool = None
config = yaml.safe_load(open("config.yaml")) 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: 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: user = into_dict(curs)
users = {} connection_pool.putconn(con)
if user_id in users: return user
return users[user_id]
return None
def new_user(message: types.Message) -> bool: def new_user(message: types.Message) -> bool:
if check_user(message.chat.id): if check_user(message.chat.id):
return False return False
update_user_settings(message.chat.id, "id", message.chat.id) con = connection_pool.getconn()
update_user_settings(message.chat.id, "type", message.chat.type) curs = con.cursor()
update_user_settings(message.chat.id, "username", message.chat.username) curs.execute(
update_user_settings(message.chat.id, "first_name", message.chat.first_name) """
update_user_settings(message.chat.id, "last_name", message.chat.last_name) INSERT INTO users (id, username, firstname, lastname, chat)
update_user_settings( VALUES (%s, %s, %s, %s, %s)
message.chat.id, "time_created", datetime.now().strftime("%Y-%m-%d %H:%M:%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) con.commit()
update_user_settings( connection_pool.putconn(con)
message.chat.id, "remind_delta", 12 * 60 * 60
) # 12 часов по умолчанию
return return
def add_note( 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: ) -> bool:
if not check_user(id): if not check_user(id):
return False return False
notes = ( con = connection_pool.getconn()
yaml.safe_load(open(user_path.format(id), "r", encoding="utf-8")) curs = con.cursor()
if os.path.exists(user_path.format(id)) curs.execute(
else {} """
SELECT id FROM lists WHERE name = %s AND user_id = %s
""",
(list, id),
) )
if list not in notes: data = curs.fetchone()
notes[list] = {"name": list, "description": "Стандартный список", "notes": []}
notes[list]["notes"].append( if not data:
{ curs.execute(
"text": text, """
"time_notif": str(time_notif) if time_notif else None, INSERT INTO lists (name, user_id)
"time_created": str(datetime.now().strftime("%Y-%m-%d %H:%M:%S")), VALUES (%s, %s)
} """,
) (list, id),
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,
) )
con.commit()
notes[list]["notes"].sort( curs.execute(
key=functools.cmp_to_key(comp), """
) 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)
yaml_data = yaml.dump( """,
notes, (id, list, id, text, remind_at),
default_flow_style=False,
encoding="utf-8",
allow_unicode=True,
width=float("inf"),
sort_keys=False,
) )
con.commit()
rowcnt = curs.rowcount
with open(user_path.format(id), "wb") as file: connection_pool.putconn(con)
file.write(yaml_data)
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): if not check_user(id):
return None return None
if not os.path.exists(user_path.format(id)): con = connection_pool.getconn()
return None 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: def delete_note(id, note_ind, list="Default") -> bool:
if not check_user(id): if not check_user(id):
return False return False
if not os.path.exists(user_path.format(id)): con = connection_pool.getconn()
return False curs = con.cursor()
curs.execute(
notes = yaml.safe_load(open(user_path.format(id), "r", encoding="utf-8")) """
DELETE FROM notes
if list not in notes: WHERE user_id = %s AND list_id = (SELECT id FROM lists WHERE name = %s AND user_id = %s)
return False ORDER BY remind_at IS NULL DESC, remind_at ASC, created_at DESC
LIMIT 1 OFFSET %s
if note_ind >= len(notes[list]["notes"]): """,
return False (id, list, id, note_ind),
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,
) )
with open(user_path.format(id), "wb") as file: con.commit()
file.write(yaml_data) rowcnt = curs.rowcount
connection_pool.putconn(con)
return True return True if rowcnt else False
def new_message(message: types.Message) -> int: def new_message(message: types.Message) -> int:
@@ -165,105 +171,98 @@ def new_message(message: types.Message) -> int:
if not check_user(id): if not check_user(id):
return None return None
users = yaml.safe_load(open(users_path, "r", encoding="utf-8")) con = connection_pool.getconn()
curs = con.cursor()
user = users[id] curs.execute(
"""
if "last_message" not in user: SELECT last_message FROM users WHERE id = %s
user["last_message"] = None """,
(id,),
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,
) )
last_id = curs.fetchone()[0]
with open(users_path, "wb") as file: curs.execute(
file.write(yaml_data) """
UPDATE users
SET last_message = %s
WHERE id = %s
""",
(new_msg_id, id),
)
con.commit()
connection_pool.putconn(con)
return last_id return last_id
def check_old_notes() -> list[int, str, int]: # возвращает первую устаревшую заметку 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: for user_id in users:
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 notes:
if not os.path.exists(user_path.format(user_id)):
continue continue
notes = yaml.safe_load(open(user_path.format(user_id), "r", encoding="utf-8"))
if "remind_delta" not in users[user_id]: timedt = timedelta(seconds=user_settings(user_id)["remind_delta"])
users[user_id]["remind_delta"] = 12 * 60 * 60 # не сохранится!
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(): return (user_id, list_name, ind)
for ind in range(len(notes[list]["notes"])):
if ( connection_pool.putconn(con)
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 None, None, None return None, None, None
def update_user_settings(user_id: int, param: str, value: any) -> bool: 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: return True if rowcnt else False
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
def user_settings(user_id: int) -> dict: def user_settings(user_id: int) -> dict:
if not check_user(user_id): con = connection_pool.getconn()
return None 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
+31
View File
@@ -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
);
+14 -36
View File
@@ -11,7 +11,7 @@ config = yaml.safe_load(open("config.yaml"))
lang = yaml.safe_load(open("lang.yaml", encoding="utf-8"))["ru"] lang = yaml.safe_load(open("lang.yaml", encoding="utf-8"))["ru"]
list_reactions = yaml.safe_load(open("lang.yaml", encoding="utf-8"))["list_reaction"] list_reactions = yaml.safe_load(open("lang.yaml", encoding="utf-8"))["list_reaction"]
bot = telebot.TeleBot( bot = telebot.TeleBot(
config["api_token"], config["test_token"],
colorful_logs=True, colorful_logs=True,
disable_web_page_preview=True, disable_web_page_preview=True,
parse_mode=ParseMode.HTML, parse_mode=ParseMode.HTML,
@@ -20,7 +20,6 @@ bot = telebot.TeleBot(
admin_id = config["admin_id"] admin_id = config["admin_id"]
bot_username = config["bot_username"] bot_username = config["bot_username"]
locale.setlocale(locale.LC_ALL, "ru_RU.UTF-8") 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]) note_ind = int(call.data.split("#")[2])
prefix = call.data.split("#")[0] 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) db.delete_note(call.message.chat.id, note_ind, prefix)
list_notes(call.message, prefix, edit=True) 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 return
@@ -92,7 +91,7 @@ def delete_notes_callback(call):
def edit_notes_callback(call): def edit_notes_callback(call):
note_ind = int(call.data.split("#")[2]) note_ind = int(call.data.split("#")[2])
prefix = call.data.split("#")[0] 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] note = notes[note_ind]
@@ -112,10 +111,10 @@ def edit_notes_callback(call):
text = f"📝 <b>Заметка {note_ind + 1}</b>" text = f"📝 <b>Заметка {note_ind + 1}</b>"
if note["time_notif"]: if note["remind_at"]:
text += f" ⌚️ <b>{datetime.strptime(note['time_notif'], '%Y-%m-%d %H:%M:%S').strftime('%d %B %Y')}</b>" text += f" ⌚️ <b>{note['remind_at'].strftime('%d %B %Y')}</b>"
text += f"\n\n{note['text']}" text += f"\n\n{note['content']}"
bot.edit_message_text( bot.edit_message_text(
text, text,
@@ -133,7 +132,7 @@ def choose_note_callback(call):
prefix = call.data.split("#")[0] prefix = call.data.split("#")[0]
markup = types.InlineKeyboardMarkup() markup = types.InlineKeyboardMarkup()
num_on_page = 6 # Лимит телеграмма - 8 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 max_page = (len(notes) - 1) // num_on_page
lst_button = [] lst_button = []
@@ -256,25 +255,6 @@ def display_settings(message: types.Message):
return return
settings = db.user_settings(message.chat.id)
# if not settings:
# send_message(message, lang["not_registered"])
# return
settings_msg = "⚙️ <u>Настройки</u>:\
\n\nВремя до отправки уведомления: <b>{}</b>"
markup = types.InlineKeyboardMarkup()
markup.add(
types.InlineKeyboardButton(
"Изменить время",
callback_data="edit_time#",
)
)
return
@bot.message_handler(commands=["list", "lists", "l"]) @bot.message_handler(commands=["list", "lists", "l"])
def list_notes(message: types.Message, list: str = "Default", edit: bool = False): 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 return
list_notes = notes["notes"]
notes_msg = "⚡️ <u>Ваши напоминания</u>:\n\n" notes_msg = "⚡️ <u>Ваши напоминания</u>:\n\n"
for i in range(len(list_notes)): for i in range(len(notes)):
notes_msg += f"<b>{i + 1})</b> " notes_msg += f"<b>{i + 1})</b> "
if list_notes[i]["time_notif"]: if notes[i]["remind_at"]:
notes_msg += f" ⌚️ <b>{datetime.strptime(list_notes[i]['time_notif'], '%Y-%m-%d %H:%M:%S').strftime('%d %B %Y')}:</b>\n" notes_msg += f" ⌚️ <b>{notes[i]['remind_at'].strftime('%d %B %Y')}:</b>\n"
notes_msg += f"{list_notes[i]['text']}\n\n" notes_msg += f"{notes[i]['content']}\n\n"
markup = types.InlineKeyboardMarkup() markup = types.InlineKeyboardMarkup()
markup.add( markup.add(
@@ -467,13 +445,13 @@ if __name__ == "__main__":
chat_id, list, ind_note = db.check_old_notes() chat_id, list, ind_note = db.check_old_notes()
if chat_id: 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) db.delete_note(chat_id, ind_note, list)
msg = bot.send_message( msg = bot.send_message(
chat_id, 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, timeout=1,
) )