74 lines
3.1 KiB
Python
74 lines
3.1 KiB
Python
import database
|
|
|
|
|
|
class VkDatabase(database.BasicDatabase):
|
|
def __init__(self, hostname: str, user: str, password: str, dbname: str):
|
|
super().__init__(hostname, user, password, dbname)
|
|
|
|
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,
|
|
api_token VARCHAR(256) NOT NULL,
|
|
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)
|
|
)""")
|
|
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS chats (
|
|
bot_id BIGINT NOT NULL,
|
|
chat_id BIGINT NOT NULL,
|
|
active TINYINT NOT NULL DEFAULT 0,
|
|
rules VARCHAR(4000),
|
|
greeting_join VARCHAR(2000),
|
|
greeting_rejoin VARCHAR(2000),
|
|
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 UPDATE CASCADE ON DELETE CASCADE
|
|
)""")
|
|
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
bot_id BIGINT NOT NULL,
|
|
chat_id BIGINT NOT NULL,
|
|
user_id BIGINT NOT NULL,
|
|
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 (bot_id, chat_id, user_id),
|
|
CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id)
|
|
ON UPDATE CASCADE ON DELETE CASCADE
|
|
)""")
|
|
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS contexts (
|
|
id BIGINT NOT NULL auto_increment,
|
|
bot_id BIGINT NOT NULL,
|
|
chat_id BIGINT NOT NULL,
|
|
message_id BIGINT,
|
|
role VARCHAR(16) NOT NULL,
|
|
message MEDIUMTEXT NOT NULL,
|
|
PRIMARY KEY (id),
|
|
CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id)
|
|
ON UPDATE CASCADE ON DELETE CASCADE
|
|
)""")
|
|
|
|
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)
|
|
|
|
|
|
DB: VkDatabase
|
|
|
|
|
|
def create_database(hostname: str, user: str, password: str, dbname: str):
|
|
global DB
|
|
DB = VkDatabase(hostname, user, password, dbname)
|