Переход с pyodbc на mariadb.
This commit is contained in:
Kirill Kirilenko
parent
56ab157a96
commit
5b9f5cd1d6
6 changed files with 284 additions and 242 deletions
224
database.py
224
database.py
|
|
@ -1,55 +1,76 @@
|
|||
from datetime import datetime
|
||||
from typing import Dict, List, Optional
|
||||
from typing import Dict, List, Optional, Any, Type
|
||||
|
||||
from pyodbc import connect, SQL_CHAR, SQL_WCHAR, Row
|
||||
from mariadb import create_pool
|
||||
|
||||
|
||||
class BasicDatabase:
|
||||
def __init__(self, connection_string: str):
|
||||
self.conn = connect(connection_string, autocommit=True)
|
||||
self.conn.setdecoding(SQL_CHAR, encoding='utf-8')
|
||||
self.conn.setdecoding(SQL_WCHAR, encoding='utf-8')
|
||||
self.conn.setencoding(encoding='utf-8')
|
||||
self.cursor = self.conn.cursor()
|
||||
def __init__(self, hostname: str, user: str, password: str, dbname: str):
|
||||
self.hostname = hostname
|
||||
self.user = user
|
||||
self.password = password
|
||||
self.dbname = dbname
|
||||
self.pool = create_pool(host=hostname, user=user, password=password, database=dbname,
|
||||
autocommit=True, dictionary=True, min_size=1, max_size=5)
|
||||
|
||||
def get_bots(self):
|
||||
self.cursor.execute("SELECT * FROM bots")
|
||||
return self._to_list(self.cursor.fetchall())
|
||||
def get_bots(self) -> List[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT * FROM bots")
|
||||
return cursor.fetchall()
|
||||
|
||||
def get_bot(self, bot_id: int):
|
||||
self.cursor.execute("SELECT * FROM bots WHERE id = ?", bot_id)
|
||||
return self._to_dict(self.cursor.fetchone())
|
||||
def get_bot(self, bot_id: int) -> Optional[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT * FROM bots WHERE id = ?", (bot_id,))
|
||||
return cursor.fetchone()
|
||||
|
||||
def get_chats(self, bot_id: int):
|
||||
self.cursor.execute("SELECT * FROM chats WHERE bot_id = ?", bot_id)
|
||||
return self._to_list(self.cursor.fetchall())
|
||||
def get_chats(self, bot_id: int) -> List[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT * FROM chats WHERE bot_id = ?", (bot_id,))
|
||||
return cursor.fetchall()
|
||||
|
||||
def get_chat(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("SELECT * FROM chats WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
return self._to_dict(self.cursor.fetchone())
|
||||
def get_chat(self, bot_id: int, chat_id: int) -> Optional[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT * FROM chats WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
|
||||
return cursor.fetchone()
|
||||
|
||||
def add_chat(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("INSERT INTO chats (bot_id, chat_id) VALUES (?, ?)", bot_id, chat_id)
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("INSERT INTO chats (bot_id, chat_id) VALUES (?, ?)", (bot_id, chat_id))
|
||||
|
||||
def chat_update(self, bot_id: int, chat_id: int, **kwargs):
|
||||
self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) +
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) +
|
||||
" WHERE bot_id = ? AND chat_id = ?", list(kwargs.values()) + [bot_id, chat_id])
|
||||
|
||||
def chat_delete(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("DELETE FROM chats WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("DELETE FROM chats WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
|
||||
|
||||
def get_user(self, bot_id: int, chat_id: int, user_id: int):
|
||||
self.cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
|
||||
bot_id, chat_id, user_id)
|
||||
return self._to_dict(self.cursor.fetchone())
|
||||
def get_user(self, bot_id: int, chat_id: int, user_id: int) -> Optional[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
|
||||
(bot_id, chat_id, user_id))
|
||||
return cursor.fetchone()
|
||||
|
||||
def get_users(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
return self._to_list(self.cursor.fetchall())
|
||||
def get_users(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
|
||||
return cursor.fetchall()
|
||||
|
||||
def add_user(self, bot_id: int, chat_id: int, user_id: int):
|
||||
self.cursor.execute("INSERT INTO users (bot_id, chat_id, user_id) VALUES (?, ?, ?)",
|
||||
bot_id, chat_id, user_id)
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("INSERT INTO users (bot_id, chat_id, user_id) VALUES (?, ?, ?)",
|
||||
(bot_id, chat_id, user_id))
|
||||
|
||||
def user_set_last_message(self, bot_id: int, chat_id: int, user_id: int, last_message: int):
|
||||
self.user_update(bot_id, chat_id, user_id, last_message=last_message)
|
||||
|
|
@ -61,83 +82,107 @@ class BasicDatabase:
|
|||
self.user_increment(bot_id, chat_id, user_id, ['warnings'])
|
||||
|
||||
def user_increment(self, bot_id: int, chat_id: int, user_id: int, fields: List[str]):
|
||||
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) +
|
||||
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?", bot_id, chat_id, user_id)
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) +
|
||||
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?", (bot_id, chat_id, user_id))
|
||||
|
||||
def user_update(self, bot_id, chat_id: int, user_id: int, **kwargs):
|
||||
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) +
|
||||
def user_update(self, bot_id: int, chat_id: int, user_id: int, **kwargs):
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) +
|
||||
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
|
||||
list(kwargs.values()) + [bot_id, chat_id, user_id])
|
||||
|
||||
def delete_user(self, bot_id: int, chat_id: int, user_id: int):
|
||||
self.cursor.execute("DELETE FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
|
||||
bot_id, chat_id, user_id)
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("DELETE FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
|
||||
(bot_id, chat_id, user_id))
|
||||
|
||||
def get_top_messages_today(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("""
|
||||
def get_top_messages_today(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("""
|
||||
SELECT user_id, messages_today AS value FROM users
|
||||
WHERE bot_id = ? AND chat_id = ? AND messages_today > 0
|
||||
ORDER BY messages_today DESC
|
||||
""", bot_id, chat_id)
|
||||
return self._to_list(self.cursor.fetchall())
|
||||
""", (bot_id, chat_id))
|
||||
return cursor.fetchall()
|
||||
|
||||
def get_top_messages_month(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("""
|
||||
def get_top_messages_month(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("""
|
||||
SELECT user_id, messages_month AS value FROM users
|
||||
WHERE bot_id = ? AND chat_id = ? AND messages_month > 0
|
||||
ORDER BY messages_month DESC
|
||||
""", bot_id, chat_id)
|
||||
return self._to_list(self.cursor.fetchall())
|
||||
""", (bot_id, chat_id))
|
||||
return cursor.fetchall()
|
||||
|
||||
def get_top_silent(self, bot_id: int, chat_id: int, threshold_days: int):
|
||||
def get_top_silent(self, bot_id: int, chat_id: int, threshold_days: int) -> List[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
current_time = int(datetime.now().timestamp())
|
||||
threshold = current_time - threshold_days * 86400
|
||||
self.cursor.execute("""
|
||||
cursor.execute("""
|
||||
SELECT user_id, (? - last_message) DIV 86400 as value
|
||||
FROM users
|
||||
WHERE bot_id = ? AND chat_id = ? AND last_message <= ?
|
||||
ORDER BY last_message ASC
|
||||
""", current_time, bot_id, chat_id, threshold)
|
||||
result = self._to_list(self.cursor.fetchall())
|
||||
""", (current_time, bot_id, chat_id, threshold))
|
||||
result = cursor.fetchall()
|
||||
for row in result:
|
||||
if row['value'] > 3650:
|
||||
row['value'] = 'никогда'
|
||||
return result
|
||||
|
||||
def get_top_warnings(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("""
|
||||
def get_top_warnings(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("""
|
||||
SELECT user_id, warnings AS value FROM users
|
||||
WHERE bot_id = ? AND chat_id = ? AND warnings > 0
|
||||
ORDER BY warnings DESC
|
||||
""", bot_id, chat_id)
|
||||
return self._to_list(self.cursor.fetchall())
|
||||
""", (bot_id, chat_id))
|
||||
return cursor.fetchall()
|
||||
|
||||
def reset_messages_today(self, bot_id: int):
|
||||
self.cursor.execute("UPDATE users SET messages_today = 0 WHERE bot_id = ?", bot_id)
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("UPDATE users SET messages_today = 0 WHERE bot_id = ?", (bot_id,))
|
||||
|
||||
def reset_messages_month(self, bot_id: int):
|
||||
self.cursor.execute("UPDATE users SET messages_month = 0 WHERE bot_id = ?", bot_id)
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("UPDATE users SET messages_month = 0 WHERE bot_id = ?", (bot_id,))
|
||||
|
||||
def context_get_messages(self, bot_id: int, chat_id: int) -> List[Dict]:
|
||||
self.cursor.execute("""
|
||||
def context_get_messages(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("""
|
||||
SELECT role, text, image FROM contexts
|
||||
WHERE bot_id = ? AND chat_id = ?
|
||||
ORDER BY id
|
||||
""", bot_id, chat_id)
|
||||
result = self._to_list(self.cursor.fetchall())
|
||||
return result
|
||||
""", (bot_id, chat_id))
|
||||
return cursor.fetchall()
|
||||
|
||||
def context_get_count(self, bot_id: int, chat_id: int) -> int:
|
||||
self.cursor.execute("SELECT COUNT(*) FROM contexts WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
return self.cursor.fetchval()
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT COUNT(*) FROM contexts WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
|
||||
return _to_val(int, cursor.fetchone())
|
||||
|
||||
def context_get_last_assistant_message_id(self, bot_id: int, chat_id: int) -> Optional[int]:
|
||||
return self.cursor.execute("""
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("""
|
||||
SELECT message_id FROM contexts
|
||||
WHERE bot_id = ? AND chat_id = ? AND role = 'assistant'
|
||||
ORDER BY id DESC
|
||||
LIMIT 1
|
||||
""", bot_id, chat_id).fetchval()
|
||||
""", (bot_id, chat_id))
|
||||
return _to_val(int, cursor.fetchone())
|
||||
|
||||
def context_add_message(self, bot_id: int, chat_id: int, role: str,
|
||||
text: Optional[str], image: Optional[bytes],
|
||||
|
|
@ -158,60 +203,55 @@ class BasicDatabase:
|
|||
placeholders = ', '.join(['?' for _ in columns])
|
||||
values = tuple(data[k] for k in columns)
|
||||
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
query = f"INSERT INTO contexts ({', '.join(columns)}) VALUES ({placeholders})"
|
||||
self.cursor.execute(query, values)
|
||||
cursor.execute(query, values)
|
||||
|
||||
def context_set_last_message_id(self, bot_id: int, chat_id: int, message_id: int):
|
||||
self.cursor.execute(
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"UPDATE contexts SET message_id = ? WHERE bot_id = ? AND chat_id = ? AND message_id IS NULL",
|
||||
message_id, bot_id, chat_id)
|
||||
(message_id, bot_id, chat_id))
|
||||
|
||||
def _context_trim(self, bot_id: int, chat_id: int, max_messages: int):
|
||||
current_count = self.context_get_count(bot_id, chat_id)
|
||||
while current_count >= max_messages:
|
||||
oldest_message_id = self.cursor.execute(
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"SELECT id FROM contexts WHERE bot_id = ? AND chat_id = ? ORDER BY id LIMIT 1",
|
||||
bot_id, chat_id).fetchval()
|
||||
(bot_id, chat_id))
|
||||
oldest_message_id = _to_val(int, cursor.fetchone())
|
||||
|
||||
if oldest_message_id:
|
||||
self.cursor.execute("DELETE FROM contexts WHERE id = ?", oldest_message_id)
|
||||
cursor.execute("DELETE FROM contexts WHERE id = ?", oldest_message_id)
|
||||
current_count -= 1
|
||||
else:
|
||||
break
|
||||
|
||||
def context_clear(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("DELETE FROM contexts WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("DELETE FROM contexts WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
|
||||
|
||||
def create_chat_if_not_exists(self, bot_id: int, chat_id: int):
|
||||
def create_chat_if_not_exists(self, bot_id: int, chat_id: int) -> Dict[str, Any]:
|
||||
chat = self.get_chat(bot_id, chat_id)
|
||||
if chat is None:
|
||||
self.add_chat(bot_id, chat_id)
|
||||
chat = self.get_chat(bot_id, chat_id)
|
||||
assert (chat is not None)
|
||||
return chat
|
||||
|
||||
def create_user_if_not_exists(self, bot_id: int, chat_id: int, user_id: int):
|
||||
def create_user_if_not_exists(self, bot_id: int, chat_id: int, user_id: int) -> Dict[str, Any]:
|
||||
user = self.get_user(bot_id, chat_id, user_id)
|
||||
if user is None:
|
||||
self.add_user(bot_id, chat_id, user_id)
|
||||
user = self.get_user(bot_id, chat_id, user_id)
|
||||
assert (user is not None)
|
||||
return user
|
||||
|
||||
def _to_dict(self, args: Optional[Row]) -> Optional[Dict]:
|
||||
columns = [column[0] for column in self.cursor.description]
|
||||
if args is not None:
|
||||
result = {}
|
||||
for i, column in enumerate(columns):
|
||||
result[column] = args[i]
|
||||
return result
|
||||
else:
|
||||
return None
|
||||
|
||||
def _to_list(self, args: List[Row]) -> List[Dict]:
|
||||
columns = [column[0] for column in self.cursor.description]
|
||||
results: list[dict] = []
|
||||
for row in args:
|
||||
row_dict = {}
|
||||
for i, column in enumerate(columns):
|
||||
row_dict[column] = row[i]
|
||||
results.append(row_dict)
|
||||
return results
|
||||
def _to_val[T](_: Type[T], arg: Any) -> T:
|
||||
return next(iter(arg.values())) if arg is not None else None
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ aiogram~=3.28.2
|
|||
aiohttp~=3.13.5
|
||||
vkbottle~=4.8.2
|
||||
vkbottle-types~=5.199.99.20
|
||||
pyodbc~=5.3.0
|
||||
mariadb[pool]~=2.0.0rc2
|
||||
openrouter==0.9.1
|
||||
replicate~=1.0.7
|
||||
tavily~=1.1.0
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ async def main() -> None:
|
|||
config = json.load(file)
|
||||
print('Конфигурация загружена.')
|
||||
|
||||
database.create_database(config['db_connection_string'])
|
||||
database.create_database(config['db_hostname'], config['db_user'], config['db_password'], config['db_database'])
|
||||
|
||||
create_ai_agent(config['openrouter_token'], config['openrouter_model'],
|
||||
config['replicate_token'], config['tavily_token'],
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@ import database
|
|||
|
||||
|
||||
class TgDatabase(database.BasicDatabase):
|
||||
def __init__(self, connection_string: str):
|
||||
super().__init__(connection_string)
|
||||
def __init__(self, hostname: str, user: str, password: str, dbname: str):
|
||||
super().__init__(hostname, user, password, dbname)
|
||||
|
||||
self.cursor.execute("""
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS bots (
|
||||
id BIGINT NOT NULL,
|
||||
owner_id BIGINT NOT NULL,
|
||||
|
|
@ -13,10 +15,10 @@ class TgDatabase(database.BasicDatabase):
|
|||
ai_prompt VARCHAR(2000) DEFAULT NULL,
|
||||
group_chats_allowed TINYINT NOT NULL DEFAULT 1,
|
||||
private_chats_allowed TINYINT NOT NULL DEFAULT 1,
|
||||
PRIMARY KEY (id))
|
||||
""")
|
||||
PRIMARY KEY (id)
|
||||
)""")
|
||||
|
||||
self.cursor.execute("""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
|
|
@ -25,10 +27,11 @@ class TgDatabase(database.BasicDatabase):
|
|||
greeting_join VARCHAR(2000),
|
||||
ai_prompt VARCHAR(2000),
|
||||
PRIMARY KEY (bot_id, chat_id),
|
||||
CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) ON DELETE CASCADE ON UPDATE CASCADE)
|
||||
""")
|
||||
CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id)
|
||||
ON UPDATE CASCADE ON DELETE CASCADE
|
||||
)""")
|
||||
|
||||
self.cursor.execute("""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
|
|
@ -43,7 +46,7 @@ class TgDatabase(database.BasicDatabase):
|
|||
ON UPDATE CASCADE ON DELETE CASCADE
|
||||
)""")
|
||||
|
||||
self.cursor.execute("""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS contexts (
|
||||
id BIGINT NOT NULL auto_increment,
|
||||
bot_id BIGINT NOT NULL,
|
||||
|
|
@ -57,12 +60,10 @@ class TgDatabase(database.BasicDatabase):
|
|||
ON UPDATE CASCADE ON DELETE CASCADE
|
||||
)""")
|
||||
|
||||
self.conn.commit()
|
||||
|
||||
|
||||
DB: TgDatabase
|
||||
|
||||
|
||||
def create_database(connection_string: str):
|
||||
def create_database(hostname: str, user: str, password: str, dbname: str):
|
||||
global DB
|
||||
DB = TgDatabase(connection_string)
|
||||
DB = TgDatabase(hostname, user, password, dbname)
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ if __name__ == '__main__':
|
|||
config = json.load(file)
|
||||
print('Конфигурация загружена.')
|
||||
|
||||
database.create_database(config['db_connection_string'])
|
||||
database.create_database(config['db_hostname'], config['db_user'], config['db_password'], config['db_database'])
|
||||
|
||||
create_ai_agent(config['openrouter_token'], config['openrouter_model'],
|
||||
config['replicate_token'], config['tavily_token'],
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@ import database
|
|||
|
||||
|
||||
class VkDatabase(database.BasicDatabase):
|
||||
def __init__(self, connection_string: str):
|
||||
super().__init__(connection_string)
|
||||
def __init__(self, hostname: str, user: str, password: str, dbname: str):
|
||||
super().__init__(hostname, user, password, dbname)
|
||||
|
||||
self.cursor.execute("""
|
||||
with self.pool.acquire() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS bots (
|
||||
id BIGINT NOT NULL,
|
||||
owner_id BIGINT NOT NULL,
|
||||
|
|
@ -13,10 +15,10 @@ class VkDatabase(database.BasicDatabase):
|
|||
ai_prompt VARCHAR(2000) DEFAULT NULL,
|
||||
group_chats_allowed TINYINT NOT NULL DEFAULT 1,
|
||||
private_chats_allowed TINYINT NOT NULL DEFAULT 1,
|
||||
PRIMARY KEY (id))
|
||||
""")
|
||||
PRIMARY KEY (id)
|
||||
)""")
|
||||
|
||||
self.cursor.execute("""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
|
|
@ -27,10 +29,11 @@ class VkDatabase(database.BasicDatabase):
|
|||
birthday_message VARCHAR(2000),
|
||||
ai_prompt VARCHAR(2000),
|
||||
PRIMARY KEY (bot_id, chat_id),
|
||||
CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) ON DELETE CASCADE ON UPDATE CASCADE)
|
||||
""")
|
||||
CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id)
|
||||
ON UPDATE CASCADE ON DELETE CASCADE
|
||||
)""")
|
||||
|
||||
self.cursor.execute("""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
|
|
@ -46,7 +49,7 @@ class VkDatabase(database.BasicDatabase):
|
|||
ON UPDATE CASCADE ON DELETE CASCADE
|
||||
)""")
|
||||
|
||||
self.cursor.execute("""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS contexts (
|
||||
id BIGINT NOT NULL auto_increment,
|
||||
bot_id BIGINT NOT NULL,
|
||||
|
|
@ -60,8 +63,6 @@ class VkDatabase(database.BasicDatabase):
|
|||
ON UPDATE CASCADE ON DELETE CASCADE
|
||||
)""")
|
||||
|
||||
self.conn.commit()
|
||||
|
||||
def user_toggle_happy_birthday(self, bot_id: int, chat_id: int, user_id: int, happy_birthday: int):
|
||||
self.user_update(bot_id, chat_id, user_id, happy_birthday=happy_birthday)
|
||||
|
||||
|
|
@ -69,6 +70,6 @@ class VkDatabase(database.BasicDatabase):
|
|||
DB: VkDatabase
|
||||
|
||||
|
||||
def create_database(connection_string: str):
|
||||
def create_database(hostname: str, user: str, password: str, dbname: str):
|
||||
global DB
|
||||
DB = VkDatabase(connection_string)
|
||||
DB = VkDatabase(hostname, user, password, dbname)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue