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
+1
View File
@@ -2,3 +2,4 @@ config.yaml
database/*
__pycache__/*
old.py
test.py
+185 -186
View File
@@ -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
+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"]
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"📝 <b>Заметка {note_ind + 1}</b>"
if note["time_notif"]:
text += f" ⌚️ <b>{datetime.strptime(note['time_notif'], '%Y-%m-%d %H:%M:%S').strftime('%d %B %Y')}</b>"
if note["remind_at"]:
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(
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 = "⚙️ <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"])
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 = "⚡️ <u>Ваши напоминания</u>:\n\n"
for i in range(len(list_notes)):
for i in range(len(notes)):
notes_msg += f"<b>{i + 1})</b> "
if list_notes[i]["time_notif"]:
notes_msg += f" ⌚️ <b>{datetime.strptime(list_notes[i]['time_notif'], '%Y-%m-%d %H:%M:%S').strftime('%d %B %Y')}:</b>\n"
if notes[i]["remind_at"]:
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.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,
)