From 7607282db2c211a24b108acad8f3fcc28c47fcb8 Mon Sep 17 00:00:00 2001 From: MythEclipse Date: Sun, 31 May 2026 23:01:38 +0700 Subject: [PATCH] feat(database): add automatic migration on startup and enhance text analysis cache source --- .env.example | 4 + README.md | 2 + .../0006_fix_text_analysis_cache_source.sql | 2 + drizzle/migrations/meta/_journal.json | 7 ++ scripts/migrate-text-analysis-cache.sql | 2 +- src/app/bootstrap.ts | 8 ++ src/config.ts | 5 ++ src/database/drizzle.ts | 22 +++++ src/database/migrate.ts | 35 ++++++-- src/database/schema.ts | 4 +- src/moderation/textCacheStore.ts | 4 +- tests/config.test.ts | 17 +++- tests/database/migrate.test.ts | 80 ++++++++++++++++++- 13 files changed, 178 insertions(+), 14 deletions(-) create mode 100644 drizzle/migrations/0006_fix_text_analysis_cache_source.sql diff --git a/.env.example b/.env.example index 6e37291..c18f620 100644 --- a/.env.example +++ b/.env.example @@ -88,3 +88,7 @@ AUTO_DELETE_ALLOWED_SEVERITIES=critical,high,medium # Optional: comma-separated category filter (empty = all categories) # AUTO_DELETE_ALLOWED_CATEGORIES= +# Database Migration Configuration +# Safe default: run migrations on startup before the app accepts traffic. +AUTO_MIGRATE_ON_STARTUP=true + diff --git a/README.md b/README.md index 38c875f..875316c 100644 --- a/README.md +++ b/README.md @@ -94,6 +94,8 @@ pnpm run install:yt-dlp Default database adalah SQLite di `.muxer-queue.db`. PostgreSQL dapat dipakai dengan `DATABASE_TYPE=postgres` dan konfigurasi `DATABASE_URL` atau variabel `POSTGRES_*`. +Jika aplikasi dijalankan di PostgreSQL production, migrasi dapat dijalankan otomatis saat startup dengan `AUTO_MIGRATE_ON_STARTUP=true`. Jalur ini memakai advisory lock supaya hanya satu instance yang memigrasi schema pada satu waktu. + ```bash # Generate migration Drizzle pnpm run db:generate diff --git a/drizzle/migrations/0006_fix_text_analysis_cache_source.sql b/drizzle/migrations/0006_fix_text_analysis_cache_source.sql new file mode 100644 index 0000000..e3276b5 --- /dev/null +++ b/drizzle/migrations/0006_fix_text_analysis_cache_source.sql @@ -0,0 +1,2 @@ +ALTER TABLE "text_analysis_cache" DROP CONSTRAINT IF EXISTS "text_analysis_cache_source_check";--> statement-breakpoint +ALTER TABLE "text_analysis_cache" ADD CONSTRAINT "text_analysis_cache_source_check" CHECK ("source" IN ('local', 'nvidia', 'primary_ai', 'groq', 'vision_llm'));--> statement-breakpoint \ No newline at end of file diff --git a/drizzle/migrations/meta/_journal.json b/drizzle/migrations/meta/_journal.json index 47a085b..a5c7297 100644 --- a/drizzle/migrations/meta/_journal.json +++ b/drizzle/migrations/meta/_journal.json @@ -36,6 +36,13 @@ "when": 1780218363790, "tag": "0005_optimize-message-index", "breakpoints": true + }, + { + "idx": 6, + "version": "7", + "when": 1780218363791, + "tag": "0006_fix_text_analysis_cache_source", + "breakpoints": true } ] } \ No newline at end of file diff --git a/scripts/migrate-text-analysis-cache.sql b/scripts/migrate-text-analysis-cache.sql index 1ab9f18..8d17065 100644 --- a/scripts/migrate-text-analysis-cache.sql +++ b/scripts/migrate-text-analysis-cache.sql @@ -5,7 +5,7 @@ CREATE TABLE IF NOT EXISTS text_analysis_cache ( text TEXT PRIMARY KEY, flags TEXT NOT NULL DEFAULT '[]', source TEXT NOT NULL DEFAULT 'local' - CHECK (source IN ('local', 'nvidia', 'primary_ai', 'groq')), + CHECK (source IN ('local', 'nvidia', 'primary_ai', 'groq', 'vision_llm')), analyzed_at BIGINT NOT NULL, expires_at BIGINT NOT NULL, hit_count INTEGER NOT NULL DEFAULT 0 diff --git a/src/app/bootstrap.ts b/src/app/bootstrap.ts index 323e409..3872a50 100644 --- a/src/app/bootstrap.ts +++ b/src/app/bootstrap.ts @@ -1,6 +1,7 @@ import { Client } from "discord.js-selfbot-v13"; import { config } from "../config.js"; import { closeDatabase, initializeDatabase } from "../database/drizzle.js"; +import { runMigrations } from "../database/migrate.js"; import { createDiscordClientOptions } from "../discordClientOptions.js"; import { createChildLogger } from "../logger.js"; import { startPendingAIAnalysisWorker } from "../moderation/aiAnalyzer.js"; @@ -40,6 +41,13 @@ export async function initializeApp() { }); try { + if (config.AUTO_MIGRATE_ON_STARTUP) { + logger.info( + "AUTO_MIGRATE_ON_STARTUP enabled; running database migrations", + ); + await runMigrations(); + } + logger.info("Initializing database"); await initializeDatabase(); logger.info("PostgreSQL database initialized"); diff --git a/src/config.ts b/src/config.ts index 0beda09..1481e79 100644 --- a/src/config.ts +++ b/src/config.ts @@ -167,6 +167,11 @@ const configSchema = z .optional() .transform((v) => v === "true") .default(true), + AUTO_MIGRATE_ON_STARTUP: z + .string() + .optional() + .transform((v) => v === "true") + .default(true), DATABASE_URL: z.string().optional(), POSTGRES_HOST: z.string().default("localhost"), POSTGRES_PORT: z.coerce.number().int().positive().default(5432), diff --git a/src/database/drizzle.ts b/src/database/drizzle.ts index 60f3bd0..079a719 100644 --- a/src/database/drizzle.ts +++ b/src/database/drizzle.ts @@ -1,5 +1,6 @@ import { drizzle as drizzlePostgres } from "drizzle-orm/node-postgres"; import { Pool } from "pg"; +import type { PoolClient } from "pg"; import { config } from "../config.js"; import { createChildLogger } from "../logger.js"; import * as schema from "./schema.js"; @@ -93,6 +94,27 @@ export async function executeGet(sql: string, params?: unknown[]) { return result.rows[0] ?? null; } +/** + * Run a function with a dedicated PostgreSQL client from the shared pool. + * Use this for session-scoped operations such as advisory locks. + */ +export async function withDatabaseClient( + callback: (client: PoolClient) => Promise, +): Promise { + if (!rawPool) { + throw new Error( + "Database not initialized. Call initializeDatabase() first.", + ); + } + + const client = await rawPool.connect(); + try { + return await callback(client); + } finally { + client.release(); + } +} + /** * Close the PostgreSQL connection pool. */ diff --git a/src/database/migrate.ts b/src/database/migrate.ts index f94a39b..4a883ca 100644 --- a/src/database/migrate.ts +++ b/src/database/migrate.ts @@ -1,19 +1,44 @@ import "dotenv/config"; +import { drizzle as drizzlePostgres } from "drizzle-orm/node-postgres"; import { migrate as migratePostgres } from "drizzle-orm/node-postgres/migrator"; +import { config } from "../config.js"; import { createChildLogger } from "../logger.js"; -import { closeDatabase, initializeDatabase } from "./drizzle.js"; +import * as schema from "./schema.js"; +import { + closeDatabase, + initializeDatabase, + withDatabaseClient, +} from "./drizzle.js"; const logger = createChildLogger("migrate"); +const MIGRATION_LOCK_KEY_1 = 2026; +const MIGRATION_LOCK_KEY_2 = 531; export async function runMigrations(): Promise { try { logger.info("Starting PostgreSQL migrations"); - const db = (await initializeDatabase()) as Parameters< - typeof migratePostgres - >[0]; + await initializeDatabase(); try { - await migratePostgres(db, { migrationsFolder: "./drizzle/migrations" }); + await withDatabaseClient(async (client) => { + const db = drizzlePostgres(client, { schema }); + + await client.query("SELECT pg_advisory_lock($1, $2)", [ + MIGRATION_LOCK_KEY_1, + MIGRATION_LOCK_KEY_2, + ]); + + try { + await migratePostgres(db, { + migrationsFolder: "./drizzle/migrations", + }); + } finally { + await client.query("SELECT pg_advisory_unlock($1, $2)", [ + MIGRATION_LOCK_KEY_1, + MIGRATION_LOCK_KEY_2, + ]); + } + }); } finally { await closeDatabase(); } diff --git a/src/database/schema.ts b/src/database/schema.ts index 8cda333..d122471 100644 --- a/src/database/schema.ts +++ b/src/database/schema.ts @@ -369,9 +369,9 @@ export const pgTextAnalysisCacheTable = pgTable( text: pgText("text").primaryKey(), /** JSON array of moderation flags detected for this text (e.g. ["vulgar_language","harassment"]). */ flags: pgText("flags").notNull().default("[]"), - /** Which source produced this result: "local" | "nvidia" | "primary_ai" | "groq". */ + /** Which source produced this result: "local" | "nvidia" | "primary_ai" | "groq" | "vision_llm". */ source: pgText("source", { - enum: ["local", "nvidia", "primary_ai", "groq"], + enum: ["local", "nvidia", "primary_ai", "groq", "vision_llm"], }) .notNull() .default("local"), diff --git a/src/moderation/textCacheStore.ts b/src/moderation/textCacheStore.ts index e1c0046..dbce397 100644 --- a/src/moderation/textCacheStore.ts +++ b/src/moderation/textCacheStore.ts @@ -7,7 +7,7 @@ const logger = createChildLogger("text-cache-store"); export interface TextCacheEntry { text: string; flags: string[]; - source: "local" | "nvidia" | "primary_ai" | "groq"; + source: "local" | "nvidia" | "primary_ai" | "groq" | "vision_llm"; analyzed_at: number; expires_at: number; hit_count: number; @@ -53,7 +53,7 @@ export async function getCachedText( export async function upsertCachedText( text: string, flags: string[], - source: "local" | "nvidia" | "primary_ai" | "groq", + source: "local" | "nvidia" | "primary_ai" | "groq" | "vision_llm", expiresAt: number, ): Promise { const now = Date.now(); diff --git a/tests/config.test.ts b/tests/config.test.ts index 591425d..6ab15eb 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -39,9 +39,10 @@ describe("loadConfig", () => { expect(config.AI_ANALYSIS_DEBOUNCE_MS).toBe(500); expect(config.AI_ANALYSIS_RECOVERY_INTERVAL_MS).toBe(15000); expect(config.AI_ANALYSIS_ERROR_COOLDOWN_MS).toBe(30000); - expect(config.AI_ANALYSIS_MAX_BATCH_SIZE).toBe(25); + expect(config.AI_ANALYSIS_MAX_BATCH_SIZE).toBe(200); expect(config.AI_ANALYSIS_MAX_CONTEXT_TOKENS).toBe(8000); expect(config.AI_ANALYSIS_CONTEXT_MESSAGE_LIMIT).toBe(20); + expect(config.AUTO_MIGRATE_ON_STARTUP).toBe(true); }); it("coerces AI analysis tuning values", async () => { @@ -107,4 +108,18 @@ describe("loadConfig", () => { expect(config.TEXT_CHANNEL_ID).toBe("text-channel"); expect(config.EFFECTIVE_VOICE_GUILD_ID).toBe("voice-guild"); }); + + it("allows disabling startup migrations explicitly", async () => { + process.env = { + ...originalEnv, + DISCORD_TOKEN: "token", + AUTO_MIGRATE_ON_STARTUP: "false", + NODE_ENV: "test", + }; + + const { loadConfig } = await import("../src/config"); + const config = loadConfig(process.env); + + expect(config.AUTO_MIGRATE_ON_STARTUP).toBe(false); + }); }); diff --git a/tests/database/migrate.test.ts b/tests/database/migrate.test.ts index ac87144..28eace4 100644 --- a/tests/database/migrate.test.ts +++ b/tests/database/migrate.test.ts @@ -1,8 +1,82 @@ -import { describe, expect, it } from "vitest"; -import { runMigrations } from "../../src/database/migrate"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const migrateSpy = vi.fn(); +const initializeDatabaseSpy = vi.fn(); +const closeDatabaseSpy = vi.fn(); +const withDatabaseClientSpy = vi.fn(); +const lockQuerySpy = vi.fn(); +const unlockQuerySpy = vi.fn(); +const drizzleSpy = vi.fn(); + +vi.mock("drizzle-orm/node-postgres", () => ({ + drizzle: drizzleSpy, +})); + +vi.mock("drizzle-orm/node-postgres/migrator", () => ({ + migrate: migrateSpy, +})); + +vi.mock("../../src/config", () => ({ + config: { + AUTO_MIGRATE_ON_STARTUP: true, + }, +})); + +vi.mock("../../src/database/drizzle", () => ({ + initializeDatabase: initializeDatabaseSpy, + closeDatabase: closeDatabaseSpy, + withDatabaseClient: withDatabaseClientSpy, +})); describe("runMigrations", () => { - it("exports the PostgreSQL migration runner", () => { + beforeEach(() => { + migrateSpy.mockResolvedValue(undefined); + initializeDatabaseSpy.mockResolvedValue({} as never); + closeDatabaseSpy.mockResolvedValue(undefined); + drizzleSpy.mockReturnValue({} as never); + withDatabaseClientSpy.mockImplementation( + async (callback: () => Promise) => + callback({ + query: vi.fn().mockImplementation((sql: string) => { + if (sql.includes("pg_advisory_lock")) { + return lockQuerySpy(sql); + } + if (sql.includes("pg_advisory_unlock")) { + return unlockQuerySpy(sql); + } + + return Promise.resolve({}); + }), + } as never), + ); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("exports the PostgreSQL migration runner", async () => { + const { runMigrations } = await import("../../src/database/migrate"); + expect(runMigrations).toBeTypeOf("function"); }); + + it("uses a session lock around the migration batch", async () => { + const { runMigrations } = await import("../../src/database/migrate"); + await runMigrations(); + + expect(initializeDatabaseSpy).toHaveBeenCalledTimes(1); + expect(withDatabaseClientSpy).toHaveBeenCalledTimes(1); + expect(migrateSpy).toHaveBeenCalledWith( + expect.any(Object), + expect.objectContaining({ migrationsFolder: "./drizzle/migrations" }), + ); + expect(lockQuerySpy).toHaveBeenCalledWith( + "SELECT pg_advisory_lock($1, $2)", + ); + expect(unlockQuerySpy).toHaveBeenCalledWith( + "SELECT pg_advisory_unlock($1, $2)", + ); + expect(closeDatabaseSpy).toHaveBeenCalledTimes(1); + }); });