tg-chat-bot/src/db/database-manager.ts

import "dotenv/config";
import {drizzle, LibSQLDatabase} from "drizzle-orm/libsql";
import {Environment} from "../common/environment";
import {logError} from "../util/utils";
import {sql} from "drizzle-orm";

type TableInfoRow = {
    name: string;
    pk: number;
};

export class DatabaseManager {

    static db: LibSQLDatabase;
    static ready: Promise<void> = Promise.resolve();

    static init() {
        DatabaseManager.db = drizzle(Environment.DB_PATH);
        DatabaseManager.ready = DatabaseManager.ensureSchema().catch(e => {
            logError(e);
            throw e;
        });
    }

    private static async getTableInfo(tableName: string): Promise<TableInfoRow[]> {
        return DatabaseManager.db.all<TableInfoRow>(sql.raw(`PRAGMA table_info(${tableName})`)).catch((e: Error) => {
            const message = String(e?.message ?? e);
            if (!message.includes("no such table")) logError(e);
            return [];
        });
    }

    private static async ensureSchema(): Promise<void> {
        await DatabaseManager.ensureUsersTable();
        await DatabaseManager.ensureMessagesTable();
    }

    private static async ensureUsersTable(): Promise<void> {
        await DatabaseManager.db.run(sql`
            CREATE TABLE IF NOT EXISTS users
            (
                id INTEGER PRIMARY KEY NOT NULL,
                isBot INTEGER NOT NULL,
                firstName TEXT NOT NULL,
                lastName TEXT,
                userName TEXT,
                isPremium INTEGER,
                langCode TEXT,
                interfaceLanguage TEXT DEFAULT 'default',
                aiProvider TEXT,
                aiResponseLanguage TEXT DEFAULT 'ru',
                aiContextSize INTEGER,
                aiVoiceMode TEXT DEFAULT 'execute'
            )
        `);

        const columns = await DatabaseManager.getTableInfo("users");
        const columnNames = new Set(columns.map(column => column.name));

        if (!columnNames.has("langCode")) {
            await DatabaseManager.db.run(sql`ALTER TABLE users ADD COLUMN langCode TEXT`);
        }

        if (!columnNames.has("aiProvider")) {
            await DatabaseManager.db.run(sql`ALTER TABLE users ADD COLUMN aiProvider TEXT`);
        }

        if (!columnNames.has("interfaceLanguage")) {
            await DatabaseManager.db.run(sql`ALTER TABLE users ADD COLUMN interfaceLanguage TEXT DEFAULT 'default'`);
        }

        if (!columnNames.has("aiResponseLanguage")) {
            await DatabaseManager.db.run(sql`ALTER TABLE users ADD COLUMN aiResponseLanguage TEXT DEFAULT 'ru'`);
        }

        if (!columnNames.has("aiContextSize")) {
            await DatabaseManager.db.run(sql`ALTER TABLE users ADD COLUMN aiContextSize INTEGER`);
        }

        if (!columnNames.has("aiVoiceMode")) {
            await DatabaseManager.db.run(sql`ALTER TABLE users ADD COLUMN aiVoiceMode TEXT DEFAULT 'execute'`);
        }
    }

    private static async createMessagesTable(): Promise<void> {
        await DatabaseManager.db.run(sql`
            CREATE TABLE IF NOT EXISTS messages
            (
                id INTEGER NOT NULL,
                chatId INTEGER NOT NULL,
                replyToMessageId INTEGER,
                fromId INTEGER NOT NULL,
                text TEXT,
                date INTEGER NOT NULL,
                photoMaxSizeFilePath TEXT,
                attachments TEXT,
                PRIMARY KEY (chatId, id)
            )
        `);
    }

    private static async ensureMessagesTable(): Promise<void> {
        let columns = await DatabaseManager.getTableInfo("messages");

        if (!columns.length) {
            await DatabaseManager.createMessagesTable();
            return;
        }

        const hasAttachments = columns.some(column => column.name === "attachments");
        const idPk = columns.find(column => column.name === "id")?.pk ?? 0;
        const chatIdPk = columns.find(column => column.name === "chatId")?.pk ?? 0;
        const hasCompositeMessageKey = idPk > 0 && chatIdPk > 0;

        if (hasAttachments && hasCompositeMessageKey) {
            return;
        }

        await DatabaseManager.recreateMessagesTable(columns);

        columns = await DatabaseManager.getTableInfo("messages");
        if (!columns.some(column => column.name === "attachments")) {
            throw new Error("Failed to ensure messages.attachments column.");
        }
    }

    private static async recreateMessagesTable(columns: TableInfoRow[]): Promise<void> {
        const legacyTable = `messages_legacy_${Date.now()}`;
        const hasAttachments = columns.some(column => column.name === "attachments");
        const attachmentsSelect = hasAttachments ? "attachments" : "NULL AS attachments";

        await DatabaseManager.db.run(sql.raw(`ALTER TABLE messages RENAME TO ${legacyTable}`));
        await DatabaseManager.createMessagesTable();
        await DatabaseManager.db.run(sql.raw(`
            INSERT OR REPLACE INTO messages
                (id, chatId, replyToMessageId, fromId, text, date, photoMaxSizeFilePath, attachments)
            SELECT id, chatId, replyToMessageId, fromId, text, date, photoMaxSizeFilePath, ${attachmentsSelect}
            FROM ${legacyTable}
        `));
        await DatabaseManager.db.run(sql.raw(`DROP TABLE ${legacyTable}`));
    }
}