storage: persist message attachments and user AI settings

src/db/database-manager.ts

@@ -2,16 +2,141 @@ 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() {
-        try {
-            DatabaseManager.db = drizzle(Environment.DB_PATH);
-        } catch (e) {
+        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}`));
+    }
+}

src/db/database.ts

@@ -2,6 +2,8 @@ import * as fs from "fs";
 import {Environment} from "../common/environment";
 import {logError} from "../util/utils";
 import {Answers} from "../model/answers";
+import path from "node:path";
+import {KeyedAsyncLock} from "../util/async-lock";
 
 type DataJsonFile = {
     admins: number[]
@@ -10,9 +12,42 @@ type DataJsonFile = {
 
 export let jsonFile: DataJsonFile;
 
+const DEFAULT_DATA: DataJsonFile = {
+    admins: [],
+    muted: [],
+};
+
+const DEFAULT_ANSWERS: Answers = {
+    test: ["a"],
+    prefix: ["?"],
+    better: ["Better"],
+    who: [],
+    kick: [],
+    invite: [],
+    day: [],
+};
+
+const dataFileLock = new KeyedAsyncLock();
+
+function ensureDataPath(): void {
+    fs.mkdirSync(Environment.DATA_PATH, {recursive: true});
+}
+
+function readJsonFile<T>(fileName: string, defaultValue: T): T {
+    ensureDataPath();
+
+    const filePath = `${Environment.DATA_PATH}/${fileName}`;
+    if (!fs.existsSync(filePath)) {
+        fs.writeFileSync(filePath, JSON.stringify(defaultValue, null, 2));
+        return structuredClone(defaultValue);
+    }
+
+    return JSON.parse(fs.readFileSync(filePath).toString()) as T;
+}
+
 export async function readData(): Promise<void> {
     try {
-        jsonFile = JSON.parse(fs.readFileSync(`${Environment.DATA_PATH}/data.json`).toString());
+        jsonFile = readJsonFile("data.json", DEFAULT_DATA);
 
         const admins = jsonFile.admins || [];
         admins.unshift(Environment.CREATOR_ID);
@@ -28,29 +63,37 @@ export async function readData(): Promise<void> {
 }
 
 export async function saveData(): Promise<void> {
-    const adminIds: number[] = [];
-    Environment.ADMIN_IDS.forEach(id => adminIds.push(id));
-    jsonFile.admins = adminIds;
+    return dataFileLock.runExclusive("data.json", async () => {
+        ensureDataPath();
+        jsonFile ??= structuredClone(DEFAULT_DATA);
 
-    const mutedList: number[] = [];
-    Environment.MUTED_IDS.forEach(id => mutedList.push(id));
-    jsonFile.muted = mutedList;
+        const adminIds: number[] = [];
+        Environment.ADMIN_IDS.forEach(id => adminIds.push(id));
+        jsonFile.admins = adminIds;
 
-    try {
-        fs.writeFileSync(`${Environment.DATA_PATH}/data.json`, JSON.stringify(jsonFile));
-        return readData();
-    } catch (e) {
-        return Promise.reject(e);
-    }
+        const mutedList: number[] = [];
+        Environment.MUTED_IDS.forEach(id => mutedList.push(id));
+        jsonFile.muted = mutedList;
+
+        try {
+            const filePath = path.join(Environment.DATA_PATH, "data.json");
+            const tmpPath = `${filePath}.${process.pid}.${Date.now()}.tmp`;
+            fs.writeFileSync(tmpPath, JSON.stringify(jsonFile));
+            fs.renameSync(tmpPath, filePath);
+            return readData();
+        } catch (e) {
+            return Promise.reject(e);
+        }
+    });
 }
 
 export async function retrieveAnswers(): Promise<void> {
     try {
-        const json: Answers = JSON.parse(fs.readFileSync(`${Environment.DATA_PATH}/answers.json`).toString());
+        const json = readJsonFile("answers.json", DEFAULT_ANSWERS);
         Environment.setAnswers(json);
         return Promise.resolve();
     } catch (e) {
         logError(e);
         return Promise.reject(e);
     }
-}
+}

src/db/message-dao.ts

@@ -5,12 +5,14 @@ import {and, eq} from "drizzle-orm";
 import {inArray} from "drizzle-orm/sql/expressions/conditions";
 import {Dao} from "../base/dao";
 import {buildExcludedSet} from "../util/utils";
+import {StoredAttachment} from "../model/stored-attachment";
 
 export class MessageDao extends Dao<StoredMessage> {
 
     private tag: string = "MessageDao";
 
     override async getAll(): Promise<StoredMessage[]> {
+        await DatabaseManager.ready;
         const then = Date.now();
 
         const messages = await DatabaseManager.db.select().from(messagesTable);
@@ -23,6 +25,7 @@ export class MessageDao extends Dao<StoredMessage> {
     }
 
     override async getById(params: { chatId: number, id: number }): Promise<StoredMessage | null> {
+        await DatabaseManager.ready;
         const then = Date.now();
 
         const messages =
@@ -45,6 +48,7 @@ export class MessageDao extends Dao<StoredMessage> {
     }
 
     override async getByIds(params: { chatId: number, ids: number[] }): Promise<StoredMessage[]> {
+        await DatabaseManager.ready;
         const then = Date.now();
 
         const messages =
@@ -65,13 +69,16 @@ export class MessageDao extends Dao<StoredMessage> {
     }
 
     async insert(values: MessageInsert[]): Promise<true> {
+        if (!values.length) return true;
+
+        await DatabaseManager.ready;
         const then = Date.now();
         const r = await DatabaseManager.db
             .insert(messagesTable)
             .values(values)
             .onConflictDoUpdate({
-                target: messagesTable.id,
-                set: buildExcludedSet(messagesTable, ["id"])
+                target: [messagesTable.chatId, messagesTable.id],
+                set: buildExcludedSet(messagesTable, ["chatId", "id"])
             });
 
         const now = Date.now();
@@ -88,8 +95,10 @@ export class MessageDao extends Dao<StoredMessage> {
                 replyToMessageId: msg.replyToMessageId,
                 fromId: msg.fromId,
                 text: msg.text,
+                quoteText: msg.quoteText,
                 date: msg.date,
                 photoMaxSizeFilePath: msg.photoMaxSizeFilePath?.join(";"),
+                attachments: msg.attachments?.length ? JSON.stringify(msg.attachments) : undefined,
             };
         });
     }
@@ -102,9 +111,22 @@ export class MessageDao extends Dao<StoredMessage> {
                 replyToMessageId: m.replyToMessageId || undefined,
                 fromId: m.fromId,
                 text: m.text,
+                quoteText: m.quoteText,
                 date: m.date,
-                photoMaxSizeFilePath: m.photoMaxSizeFilePath?.split(";")
+                photoMaxSizeFilePath: m.photoMaxSizeFilePath?.split(";"),
+                attachments: parseAttachments(m.attachments),
             };
         });
     }
-}
+}
+
+function parseAttachments(value?: string | null): StoredAttachment[] | undefined {
+    if (!value?.trim()) return undefined;
+
+    try {
+        const parsed = JSON.parse(value);
+        return Array.isArray(parsed) ? parsed : undefined;
+    } catch {
+        return undefined;
+    }
+}

src/db/schema.ts

@@ -1,24 +1,34 @@
-import {int, sqliteTable, text} from "drizzle-orm/sqlite-core";
+import {int, primaryKey, sqliteTable, text} from "drizzle-orm/sqlite-core";
 
 export const messagesTable = sqliteTable("messages", {
-    id: int().primaryKey().unique().notNull(),
+    id: int().notNull(),
     chatId: int().notNull(),
     replyToMessageId: int(),
     fromId: int().notNull(),
     text: text(),
+    quoteText: text(),
     date: int().notNull(),
     photoMaxSizeFilePath: text(),
-});
+    attachments: text(),
+}, table => [
+    primaryKey({columns: [table.chatId, table.id]}),
+]);
 
 export type MessageInsert = typeof messagesTable.$inferInsert;
 
 export const usersTable = sqliteTable("users", {
-    id: int().primaryKey().unique().notNull(),
+    id: int().primaryKey().notNull(),
     isBot: int().notNull(),
     firstName: text().notNull(),
     lastName: text(),
     userName: text(),
     isPremium: int(),
+    langCode: text(),
+    interfaceLanguage: text().default("default"),
+    aiProvider: text(),
+    aiResponseLanguage: text().default("ru"),
+    aiContextSize: int(),
+    aiVoiceMode: text().default("execute"),
 });
 
 export type UserInsert = typeof usersTable.$inferInsert;

src/db/user-dao.ts

@@ -12,6 +12,7 @@ export class UserDao extends Dao<StoredUser> {
     private tag: string = "UserDao";
 
     override async getAll(): Promise<StoredUser[]> {
+        await DatabaseManager.ready;
         const then = Date.now();
 
         const users = await DatabaseManager.db.select().from(usersTable);
@@ -24,6 +25,7 @@ export class UserDao extends Dao<StoredUser> {
     }
 
     override async getById(params: { id: number }): Promise<StoredUser | null> {
+        await DatabaseManager.ready;
         const then = Date.now();
 
         const users =
@@ -43,6 +45,7 @@ export class UserDao extends Dao<StoredUser> {
     }
 
     override async getByIds(params: { ids: number[] }): Promise<StoredUser[]> {
+        await DatabaseManager.ready;
         const then = Date.now();
 
         const users =
@@ -60,7 +63,9 @@ export class UserDao extends Dao<StoredUser> {
     }
 
     override async insert(values: UserInsert[] | UserInsert): Promise<true> {
+        await DatabaseManager.ready;
         const rows = Array.isArray(values) ? values : [values];
+        if (!rows.length) return true;
 
         const then = Date.now();
         const r = await DatabaseManager.db
@@ -68,7 +73,7 @@ export class UserDao extends Dao<StoredUser> {
             .values(rows)
             .onConflictDoUpdate({
                 target: usersTable.id,
-                set: buildExcludedSet(usersTable, ["id"])
+                set: buildExcludedSet(usersTable, ["id", "interfaceLanguage", "aiProvider", "aiResponseLanguage", "aiContextSize", "aiVoiceMode"])
             });
 
         const now = Date.now();
@@ -77,6 +82,28 @@ export class UserDao extends Dao<StoredUser> {
         return true;
     }
 
+    async updateSettings(
+        id: number,
+        settings: Partial<Pick<StoredUser, "interfaceLanguage" | "aiProvider" | "aiResponseLanguage" | "aiContextSize" | "aiVoiceMode">>
+    ): Promise<true> {
+        await DatabaseManager.ready;
+
+        const update: Partial<UserInsert> = {};
+        if ("interfaceLanguage" in settings) update.interfaceLanguage = settings.interfaceLanguage ?? null;
+        if ("aiProvider" in settings) update.aiProvider = settings.aiProvider ?? null;
+        if ("aiResponseLanguage" in settings) update.aiResponseLanguage = settings.aiResponseLanguage ?? null;
+        if ("aiContextSize" in settings) update.aiContextSize = settings.aiContextSize ?? null;
+        if ("aiVoiceMode" in settings) update.aiVoiceMode = settings.aiVoiceMode ?? null;
+        if (!Object.keys(update).length) return true;
+
+        await DatabaseManager.db
+            .update(usersTable)
+            .set(update)
+            .where(eq(usersTable.id, id));
+
+        return true;
+    }
+
     mapTo(users: User[]): UserInsert[] {
         return users.map(u => {
             return {
@@ -85,21 +112,29 @@ export class UserDao extends Dao<StoredUser> {
                 firstName: u.first_name,
                 lastName: u.last_name,
                 userName: u.username,
-                isPremium: boolToInt(u.is_premium)
+                isPremium: boolToInt(u.is_premium),
+                langCode: u.language_code
             };
         });
     }
 
     mapFrom(users: UserInsert[]): StoredUser[] {
+        // @ts-ignore
         return users.map(u => {
             return {
                 id: <number>u.id,
                 isBot: u.isBot === 1,
                 firstName: u.firstName,
-                lastName: u.lastName,
-                userName: u.userName,
-                isPremium: u.isPremium === 1
+                lastName: u.lastName === null ? undefined : u.lastName,
+                userName: u.userName === null ? undefined : u.userName,
+                isPremium: u.isPremium === 1,
+                langCode: u.langCode === null ? undefined : u.langCode,
+                interfaceLanguage: u.interfaceLanguage === null ? undefined : u.interfaceLanguage,
+                aiProvider: u.aiProvider === null ? undefined : u.aiProvider,
+                aiResponseLanguage: u.aiResponseLanguage === null ? undefined : u.aiResponseLanguage,
+                aiContextSize: u.aiContextSize === null ? undefined : u.aiContextSize,
+                aiVoiceMode: u.aiVoiceMode === null ? undefined : u.aiVoiceMode,
             };
         });
     }
-}
+}