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,
)