Переход с SQLite3 на ODBC.
This commit is contained in:
parent
099fd5efea
commit
4942a84325
5 changed files with 104 additions and 67 deletions
81
database.py
81
database.py
|
|
@ -1,44 +1,43 @@
|
|||
import sqlite3
|
||||
from typing import List
|
||||
from pyodbc import connect, Row
|
||||
from typing import List, Union
|
||||
|
||||
|
||||
class BasicDatabase:
|
||||
def __init__(self, filename: str):
|
||||
self.conn = sqlite3.connect(filename)
|
||||
self.conn.row_factory = sqlite3.Row
|
||||
def __init__(self, connection_string: str):
|
||||
self.conn = connect(connection_string)
|
||||
self.cursor = self.conn.cursor()
|
||||
|
||||
def get_chats(self):
|
||||
self.cursor.execute("SELECT * FROM chats")
|
||||
return self.cursor.fetchall()
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def get_chat(self, chat_id: int):
|
||||
self.cursor.execute("SELECT * FROM chats WHERE id = ?", (chat_id,))
|
||||
return self.cursor.fetchone()
|
||||
self.cursor.execute("SELECT * FROM chats WHERE id = ?", chat_id)
|
||||
return self._to_dict(self.cursor.fetchone())
|
||||
|
||||
def add_chat(self, chat_id: int):
|
||||
self.cursor.execute("INSERT INTO chats (id) VALUES (?)", (chat_id,))
|
||||
self.cursor.execute("INSERT INTO chats (id) VALUES (?)", chat_id)
|
||||
self.conn.commit()
|
||||
|
||||
def chat_update(self, chat_id: int, **kwargs):
|
||||
self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = :" + f for f in kwargs) +
|
||||
f" WHERE id = {chat_id}", kwargs)
|
||||
self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) +
|
||||
" WHERE id = ?", list(kwargs.values()) + [chat_id])
|
||||
self.conn.commit()
|
||||
|
||||
def chat_delete(self, chat_id: int):
|
||||
self.cursor.execute("DELETE FROM chats WHERE id = ?", (chat_id,))
|
||||
self.cursor.execute("DELETE FROM chats WHERE id = ?", chat_id)
|
||||
self.conn.commit()
|
||||
|
||||
def get_user(self, chat_id: int, user_id: int):
|
||||
self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
|
||||
return self.cursor.fetchone()
|
||||
self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
|
||||
return self._to_dict(self.cursor.fetchone())
|
||||
|
||||
def get_users(self, chat_id: int):
|
||||
self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", (chat_id,))
|
||||
return self.cursor.fetchall()
|
||||
self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def add_user(self, chat_id: int, user_id: int):
|
||||
self.cursor.execute("INSERT INTO users (chat_id, user_id) VALUES (?, ?)", (chat_id, user_id))
|
||||
self.cursor.execute("INSERT INTO users (chat_id, user_id) VALUES (?, ?)", chat_id, user_id)
|
||||
self.conn.commit()
|
||||
|
||||
def user_set_last_message(self, chat_id: int, user_id: int, last_message: int):
|
||||
|
|
@ -53,16 +52,18 @@ class BasicDatabase:
|
|||
|
||||
def user_increment(self, chat_id: int, user_id: int, fields: List[str]):
|
||||
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) +
|
||||
" WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
|
||||
" WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
|
||||
self.conn.commit()
|
||||
|
||||
def user_update(self, chat_id: int, user_id: int, **kwargs):
|
||||
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = :" + f for f in kwargs) +
|
||||
f" WHERE chat_id = {chat_id} AND user_id = {user_id}", kwargs)
|
||||
sql = "UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) + " WHERE chat_id = ? AND user_id = ?"
|
||||
args = [chat_id, user_id] + list(kwargs.values())
|
||||
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) +
|
||||
" WHERE chat_id = ? AND user_id = ?", list(kwargs.values()) + [chat_id, user_id])
|
||||
self.conn.commit()
|
||||
|
||||
def delete_user(self, chat_id: int, user_id: int):
|
||||
self.cursor.execute("DELETE FROM users WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
|
||||
self.cursor.execute("DELETE FROM users WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
|
||||
self.conn.commit()
|
||||
|
||||
def get_top_messages_today(self, chat_id: int):
|
||||
|
|
@ -70,16 +71,16 @@ class BasicDatabase:
|
|||
SELECT user_id, messages_today AS value FROM users
|
||||
WHERE chat_id = ? AND messages_today > 0
|
||||
ORDER BY messages_today DESC
|
||||
""", (chat_id,))
|
||||
return self.cursor.fetchall()
|
||||
""", chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def get_top_messages_month(self, chat_id: int):
|
||||
self.cursor.execute("""
|
||||
SELECT user_id, messages_month AS value FROM users
|
||||
WHERE chat_id = ? AND messages_month > 0
|
||||
ORDER BY messages_month DESC
|
||||
""", (chat_id,))
|
||||
return self.cursor.fetchall()
|
||||
""", chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def get_top_silent(self, chat_id: int, threshold_days: int):
|
||||
# noinspection SpellCheckingInspection
|
||||
|
|
@ -92,19 +93,19 @@ class BasicDatabase:
|
|||
FROM users
|
||||
WHERE chat_id = ? AND value >= ?
|
||||
ORDER BY value DESC
|
||||
""", (chat_id, threshold_days))
|
||||
return self.cursor.fetchall()
|
||||
""", chat_id, threshold_days)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def get_top_warnings(self, chat_id: int):
|
||||
self.cursor.execute("""
|
||||
SELECT user_id, warnings AS value FROM users
|
||||
WHERE chat_id = ? AND warnings > 0
|
||||
ORDER BY warnings DESC
|
||||
""", (chat_id,))
|
||||
return self.cursor.fetchall()
|
||||
""", chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def delete_users(self, chat_id: int):
|
||||
self.cursor.execute("DELETE FROM users WHERE chat_id = ?", (chat_id,))
|
||||
self.cursor.execute("DELETE FROM users WHERE chat_id = ?", chat_id)
|
||||
self.conn.commit()
|
||||
|
||||
def reset_messages_today(self):
|
||||
|
|
@ -128,3 +129,23 @@ class BasicDatabase:
|
|||
self.add_user(chat_id, user_id)
|
||||
user = self.get_user(chat_id, user_id)
|
||||
return user
|
||||
|
||||
def _to_dict(self, args: Union[Row, List[Row], None]):
|
||||
columns = [column[0] for column in self.cursor.description]
|
||||
if args is None:
|
||||
return None
|
||||
elif isinstance(args, Row):
|
||||
result = {}
|
||||
for i, column in enumerate(columns):
|
||||
result[column] = args[i]
|
||||
return result
|
||||
elif isinstance(args, list) and all(isinstance(item, Row) for item in args):
|
||||
results = []
|
||||
for row in args:
|
||||
row_dict = {}
|
||||
for i, column in enumerate(columns):
|
||||
row_dict[column] = row[i]
|
||||
results.append(row_dict)
|
||||
return results
|
||||
else:
|
||||
raise TypeError("unexpected type")
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import json
|
|||
from aiogram import Bot, Dispatcher
|
||||
|
||||
from ai_agent import create_ai_agent
|
||||
from tg.tg_database import create_database
|
||||
|
||||
from . import handlers
|
||||
from . import tasks
|
||||
|
|
@ -15,6 +16,7 @@ async def main() -> None:
|
|||
print('Конфигурация загружена.')
|
||||
|
||||
bot = Bot(token=config['api_token'])
|
||||
create_database(config['db_connection_string'])
|
||||
create_ai_agent(config['openrouter_token'])
|
||||
|
||||
dp = Dispatcher()
|
||||
|
|
|
|||
|
|
@ -2,32 +2,38 @@ import database
|
|||
|
||||
|
||||
class TgDatabase(database.BasicDatabase):
|
||||
def __init__(self):
|
||||
super().__init__('tg.db')
|
||||
def __init__(self, connection_string: str):
|
||||
super().__init__(connection_string)
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
"id" INTEGER,
|
||||
"active" INTEGER NOT NULL DEFAULT 0,
|
||||
"rules" TEXT,
|
||||
"greeting_join" TEXT,
|
||||
"ai_prompt" TEXT,
|
||||
PRIMARY KEY("id"))
|
||||
id BIGINT,
|
||||
active TINYINT NOT NULL DEFAULT 0,
|
||||
rules VARCHAR(2000),
|
||||
greeting_join VARCHAR(1000),
|
||||
ai_prompt VARCHAR(2048),
|
||||
PRIMARY KEY (id))
|
||||
""")
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
"chat_id" INTEGER,
|
||||
"user_id" INTEGER,
|
||||
"last_message" INTEGER NOT NULL DEFAULT 0,
|
||||
"messages_today" INTEGER NOT NULL DEFAULT 0,
|
||||
"messages_month" INTEGER NOT NULL DEFAULT 0,
|
||||
"warnings" INTEGER NOT NULL DEFAULT 0,
|
||||
"about" TEXT,
|
||||
PRIMARY KEY("chat_id","user_id"))
|
||||
chat_id BIGINT,
|
||||
user_id BIGINT,
|
||||
last_message BIGINT NOT NULL DEFAULT 0,
|
||||
messages_today SMALLINT NOT NULL DEFAULT 0,
|
||||
messages_month SMALLINT NOT NULL DEFAULT 0,
|
||||
warnings TINYINT NOT NULL DEFAULT 0,
|
||||
about VARCHAR(1000),
|
||||
PRIMARY KEY (chat_id, user_id),
|
||||
CONSTRAINT fk_users_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON DELETE CASCADE)
|
||||
""")
|
||||
|
||||
self.conn.commit()
|
||||
|
||||
|
||||
DB = TgDatabase()
|
||||
DB: TgDatabase
|
||||
|
||||
|
||||
def create_database(connection_string: str):
|
||||
global DB
|
||||
DB = TgDatabase(connection_string)
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import json
|
|||
from vkbottle.bot import Bot as VkBot
|
||||
|
||||
from ai_agent import create_ai_agent
|
||||
from vk.vk_database import create_database
|
||||
|
||||
from . import handlers
|
||||
from . import tasks
|
||||
|
|
@ -14,6 +15,7 @@ if __name__ == '__main__':
|
|||
print('Конфигурация загружена.')
|
||||
|
||||
bot = VkBot(config['api_token'], labeler=handlers.labeler)
|
||||
create_database(config['db_connection_string'])
|
||||
create_ai_agent(config["openrouter_token"])
|
||||
|
||||
bot.loop_wrapper.on_startup.append(tasks.startup_task(bot.api))
|
||||
|
|
|
|||
|
|
@ -2,32 +2,33 @@ import database
|
|||
|
||||
|
||||
class VkDatabase(database.BasicDatabase):
|
||||
def __init__(self):
|
||||
super().__init__('vk.db')
|
||||
def __init__(self, connection_string: str):
|
||||
super().__init__(connection_string)
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
"id" INTEGER,
|
||||
"active" INTEGER NOT NULL DEFAULT 0,
|
||||
"rules" TEXT,
|
||||
"greeting_join" TEXT,
|
||||
"greeting_rejoin" TEXT,
|
||||
"birthday_message" TEXT,
|
||||
"ai_prompt" TEXT,
|
||||
PRIMARY KEY("id"))
|
||||
id BIGINT,
|
||||
active TINYINT NOT NULL DEFAULT 0,
|
||||
rules VARCHAR(2000),
|
||||
greeting_join VARCHAR(1000),
|
||||
greeting_rejoin VARCHAR(1000),
|
||||
birthday_message VARCHAR(1000),
|
||||
ai_prompt VARCHAR(2048),
|
||||
PRIMARY KEY (id))
|
||||
""")
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
"chat_id" INTEGER,
|
||||
"user_id" INTEGER,
|
||||
"last_message" INTEGER NOT NULL DEFAULT 0,
|
||||
"messages_today" INTEGER NOT NULL DEFAULT 0,
|
||||
"messages_month" INTEGER NOT NULL DEFAULT 0,
|
||||
"warnings" INTEGER NOT NULL DEFAULT 0,
|
||||
"happy_birthday" INTEGER NOT NULL DEFAULT 1,
|
||||
"about" TEXT,
|
||||
PRIMARY KEY("chat_id","user_id"))
|
||||
chat_id BIGINT,
|
||||
user_id BIGINT,
|
||||
last_message BIGINT NOT NULL DEFAULT 0,
|
||||
messages_today SMALLINT NOT NULL DEFAULT 0,
|
||||
messages_month SMALLINT NOT NULL DEFAULT 0,
|
||||
warnings TINYINT NOT NULL DEFAULT 0,
|
||||
happy_birthday TINYINT NOT NULL DEFAULT 1,
|
||||
about VARCHAR(1000),
|
||||
PRIMARY KEY (chat_id, user_id),
|
||||
CONSTRAINT fk_users_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON DELETE CASCADE)
|
||||
""")
|
||||
|
||||
self.conn.commit()
|
||||
|
|
@ -36,4 +37,9 @@ class VkDatabase(database.BasicDatabase):
|
|||
self.user_update(chat_id, user_id, happy_birthday=happy_birthday)
|
||||
|
||||
|
||||
DB = VkDatabase()
|
||||
DB: VkDatabase
|
||||
|
||||
|
||||
def create_database(connection_string: str):
|
||||
global DB
|
||||
DB = VkDatabase(connection_string)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue