ai-chatbot/lib/db/migrate.ts

import { config } from "dotenv";
import { sql } from "drizzle-orm";
import { drizzle } from "drizzle-orm/postgres-js";
import { migrate } from "drizzle-orm/postgres-js/migrator";
import { readFileSync, readdirSync } from "node:fs";
import { join } from "node:path";
import postgres from "postgres";

config({
  path: ".env.local",
});

const providerBaseUrlDefaults = [
  { type: "alibaba", baseUrl: "https://dashscope.aliyuncs.com/compatible-mode/v1" },
  { type: "openai", baseUrl: "https://api.openai.com/v1" },
  { type: "grsai-gpt-image", baseUrl: "https://grsai.dakka.com.cn" },
  { type: "grsai-video", baseUrl: "https://grsai.dakka.com.cn" },
  { type: "huanke-gemini", baseUrl: "https://api.huandutech.com/v1" },
  {
    type: "qianwen-video",
    baseUrl: "https://dashscope.aliyuncs.com/api/v1/services/aigc/video-generation",
  },
  { type: "doubao-video", baseUrl: "https://ark.cn-beijing.volces.com/api/v3" },
  { type: "ruxa-video", baseUrl: "https://api.ruxa.ai" },
  { type: "volcano-vision", baseUrl: "https://ark.cn-beijing.volces.com" },
  { type: "volcano-video", baseUrl: "https://ark.cn-beijing.volces.com/api/v3" },
  {
    type: "alibaba-image",
    baseUrl: "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation",
  },
  { type: "bfl-image", baseUrl: "https://api.bfl.ai" },
  { type: "ruxa-image", baseUrl: "https://api.ruxa.ai" },
  { type: "fal-image", baseUrl: "https://fal.run" },
  { type: "snapgen-image", baseUrl: "https://api.snapgen.ai" },
  { type: "openai-image", baseUrl: "https://api.openai.com/v1" },
  { type: "bananaimage", baseUrl: "https://bananaimage.aigenmedia.art" },
  { type: "visionary", baseUrl: "https://api.visionary.beer" },
  { type: "atlascloud", baseUrl: "https://api.atlascloud.ai" },
  { type: "vip98k", baseUrl: "https://api.vip98k.com" },
  { type: "stepfun-image", baseUrl: "https://api.stepfun.com" },
  { type: "midjourney", baseUrl: "https://your-midjourney-api.example.com" },
  { type: "fal-video", baseUrl: "https://fal.run" },
  { type: "tripo-3d", baseUrl: "https://api.tripo3d.ai/v2/openapi" },
  { type: "worldlabs-3d", baseUrl: "https://api.vip98k.com" },
  { type: "kittycad-cad", baseUrl: "https://api.zoo.dev" },
  { type: "runway", baseUrl: "https://api.dev.runwayml.com" },
  { type: "krea", baseUrl: "https://api.krea.ai" },
] as const;

function assertMigrationJournalIsConsistent() {
  const migrationsDir = join(process.cwd(), "lib/db/migrations");
  const journalPath = join(migrationsDir, "meta", "_journal.json");
  const journal = JSON.parse(readFileSync(journalPath, "utf8")) as {
    entries: { tag: string }[];
  };

  const tags = new Set(journal.entries.map((entry) => entry.tag));
  const sqlFiles = readdirSync(migrationsDir)
    .filter((file) => file.endsWith(".sql"))
    .map((file) => file.slice(0, -4));
  const sqlSet = new Set(sqlFiles);

  const extraSqlFiles = sqlFiles.filter((tag) => !tags.has(tag));
  const missingSqlFiles = Array.from(tags).filter((tag) => !sqlSet.has(tag));

  if (extraSqlFiles.length === 0 && missingSqlFiles.length === 0) {
    return;
  }

  const details = [
    extraSqlFiles.length > 0
      ? `extra sql files not in journal: ${extraSqlFiles.join(", ")}`
      : null,
    missingSqlFiles.length > 0
      ? `journal entries missing sql files: ${missingSqlFiles.join(", ")}`
      : null,
  ]
    .filter(Boolean)
    .join(" | ");

  throw new Error(
    `Migration journal mismatch. Keep lib/db/migrations/*.sql in sync with meta/_journal.json. ${details}`
  );
}

function getEnvInt(name: string, fallback: number, min: number, max: number) {
  const raw = process.env[name];
  if (!raw) {
    return fallback;
  }

  const value = Number.parseInt(raw, 10);
  if (!Number.isFinite(value)) {
    return fallback;
  }

  return Math.min(Math.max(value, min), max);
}

function withTimeout<T>(promise: Promise<T>, timeoutMs: number, message: string) {
  let timeout: NodeJS.Timeout | undefined;

  const timeoutPromise = new Promise<never>((_, reject) => {
    timeout = setTimeout(() => reject(new Error(message)), timeoutMs);
  });

  return Promise.race([promise, timeoutPromise]).finally(() => {
    if (timeout) {
      clearTimeout(timeout);
    }
  });
}

const runMigrate = async () => {
  if (!process.env.POSTGRES_URL) {
    throw new Error("POSTGRES_URL is not defined");
  }

  const migrationTimeoutSeconds = getEnvInt("POSTGRES_MIGRATION_TIMEOUT_SECONDS", 120, 5, 1800);
  const connection = postgres(process.env.POSTGRES_URL, {
    connect_timeout: getEnvInt("POSTGRES_CONNECT_TIMEOUT_SECONDS", 5, 1, 60),
    idle_timeout: getEnvInt("POSTGRES_IDLE_TIMEOUT_SECONDS", 20, 1, 300),
    max: 1,
    prepare: false,
  });
  const db = drizzle(connection);

  console.log("⏳ Running migrations...");

  const start = Date.now();
  assertMigrationJournalIsConsistent();
  try {
    await withTimeout(
      runMigrationStatements(db),
      migrationTimeoutSeconds * 1000,
      `Migration timed out after ${migrationTimeoutSeconds}s. Check POSTGRES_URL, network access, and migration locks.`
    );
  } finally {
    await connection.end({ timeout: 5 });
  }
  const end = Date.now();

  console.log("✅ Migrations completed in", end - start, "ms");
  process.exit(0);
};

async function runMigrationStatements(db: ReturnType<typeof drizzle>) {
  await migrate(db, { migrationsFolder: "./lib/db/migrations" });
  await db.execute(sql`
      ALTER TABLE "AiProvider" DROP CONSTRAINT IF EXISTS "AiProvider_type_check";
    `);
  await db.execute(sql`
      ALTER TABLE "AiProvider"
      ADD CONSTRAINT "AiProvider_type_check"
      CHECK (type IN (
        'alibaba', 'openai', 'custom', 'local', 'gemini',
        'grsai-gpt-image', 'grsai-video', 'huanke-gemini', 'qianwen-video', 'doubao-video', 'ruxa-video', 'volcano-vision', 'volcano-video', 'alibaba-image', 'bfl-image', 'ruxa-image', 'fal-image', 'snapgen-image', 'openai-image', 'bananaimage', 'visionary', 'atlascloud', 'vip98k', 'stepfun-image', 'midjourney', 'fal-video', 'zhipu', 'tripo-3d', 'worldlabs-3d', 'kittycad-cad', 'runway', 'krea'
      ));
    `);
  await db.execute(sql`
      ALTER TABLE "AiModel"
      ADD COLUMN IF NOT EXISTS "excludeFromFailover" boolean NOT NULL DEFAULT false;
    `);
  await db.execute(sql`
      ALTER TABLE "AiModel"
      ADD COLUMN IF NOT EXISTS "excludeFromHealthBan" boolean NOT NULL DEFAULT false;
    `);
  await db.execute(sql`
      ALTER TABLE "AiModel"
      ADD COLUMN IF NOT EXISTS "isDefaultTextEditModel" boolean NOT NULL DEFAULT false;
    `);
  await db.execute(sql`
      ALTER TABLE "AiModel"
      ADD COLUMN IF NOT EXISTS "healthRecoverySuccessStreak" integer NOT NULL DEFAULT 0;
    `);
  await db.execute(sql`
      ALTER TABLE "AiModel"
      ADD COLUMN IF NOT EXISTS "healthObservationStartedAt" timestamp;
    `);
  for (const item of providerBaseUrlDefaults) {
    await db.execute(sql`
      UPDATE "AiProvider"
      SET
        "baseUrl" = ${item.baseUrl},
        "updatedAt" = NOW()
      WHERE
        "type" = ${item.type}
        AND ("baseUrl" IS NULL OR BTRIM("baseUrl") = '');
    `);
  }
}

runMigrate().catch((err) => {
  console.error("❌ Migration failed");
  console.error(err);
  process.exit(1);
});