ai-chatbot/lib/storage/migration-db.ts

import "server-only";

import { and, asc, desc, eq, inArray, lte, sql } from "drizzle-orm";
import { db } from "@/lib/db";
import { storageConfig, storageMigrationTask, type StorageConfig } from "@/lib/db/schema";

export type StorageMigrationTaskRecord = typeof storageMigrationTask.$inferSelect;

export function isAliyunStorageConfig(config: Pick<StorageConfig, "type" | "endpoint"> | null | undefined) {
  if (!config) return false;
  return config.type === "s3" && /aliyuncs\.com/i.test(String(config.endpoint || ""));
}

function isUsableS3Config(config: StorageConfig | null | undefined): config is StorageConfig {
  return Boolean(
    config &&
      config.type === "s3" &&
      config.isEnabled &&
      String(config.endpoint || "").trim() &&
      String(config.bucket || "").trim() &&
      String(config.accessKeyId || "").trim() &&
      String(config.secretAccessKey || "").trim()
  );
}

export async function getEnabledS3StorageConfigs(): Promise<StorageConfig[]> {
  const configs = await db
    .select()
    .from(storageConfig)
    .where(and(eq(storageConfig.type, "s3"), eq(storageConfig.isEnabled, true)))
    .orderBy(desc(storageConfig.isDefault), asc(storageConfig.createdAt));

  return configs.filter(isUsableS3Config);
}

export async function getPreferredAliyunStorageConfig({
  preferredId,
}: {
  preferredId?: string;
} = {}): Promise<StorageConfig | undefined> {
  const configs = await getEnabledS3StorageConfigs();
  const aliyunConfigs = configs.filter((config) => isAliyunStorageConfig(config));

  if (aliyunConfigs.length === 0) return undefined;

  if (preferredId) {
    const preferred = aliyunConfigs.find((item) => item.id === preferredId);
    if (preferred) return preferred;
  }

  const defaultAliyun = aliyunConfigs.find((item) => item.isDefault);
  return defaultAliyun || aliyunConfigs[0];
}

export async function getStorageConfigByIdForMigration({
  id,
}: {
  id?: string | null;
}): Promise<StorageConfig | undefined> {
  if (!id) return undefined;
  const [config] = await db
    .select()
    .from(storageConfig)
    .where(eq(storageConfig.id, id))
    .limit(1);

  return isUsableS3Config(config) ? config : undefined;
}

export async function enqueueStorageMigrationTask({
  sourceUrl,
  sourcePath,
  sourceStorageConfigId,
  migrateAfterAt,
  cleanupAfterAt,
}: {
  sourceUrl: string;
  sourcePath: string;
  sourceStorageConfigId: string;
  migrateAfterAt: Date;
  cleanupAfterAt: Date;
}) {
  if (!sourceUrl || !sourcePath || !sourceStorageConfigId) return;

  await db
    .insert(storageMigrationTask)
    .values({
      sourceUrl,
      sourcePath,
      sourceStorageConfigId,
      migrateAfterAt,
      cleanupAfterAt,
      status: "pending",
      updatedAt: new Date(),
    })
    .onConflictDoNothing({
      target: storageMigrationTask.sourceUrl,
    });
}

export async function getStorageTasksToMigrate({
  limit,
}: {
  limit: number;
}) {
  const now = new Date();
  return await db
    .select()
    .from(storageMigrationTask)
    .where(
      and(
        inArray(storageMigrationTask.status, ["pending", "failed"]),
        lte(storageMigrationTask.migrateAfterAt, now)
      )
    )
    .orderBy(asc(storageMigrationTask.migrateAfterAt))
    .limit(Math.max(1, Math.min(100, limit)));
}

export async function markStorageTaskMigrating(id: string) {
  await db
    .update(storageMigrationTask)
    .set({
      status: "migrating",
      attemptCount: sql`${storageMigrationTask.attemptCount} + 1`,
      updatedAt: new Date(),
    })
    .where(eq(storageMigrationTask.id, id));
}

export async function markStorageTaskMigrated({
  id,
  targetUrl,
  targetStorageConfigId,
}: {
  id: string;
  targetUrl: string;
  targetStorageConfigId: string;
}) {
  await db
    .update(storageMigrationTask)
    .set({
      status: "migrated",
      targetUrl,
      targetStorageConfigId,
      migratedAt: new Date(),
      lastError: null,
      updatedAt: new Date(),
    })
    .where(eq(storageMigrationTask.id, id));
}

export async function markStorageTaskMigrationFailed({
  id,
  error,
}: {
  id: string;
  error: unknown;
}) {
  const message = String(error instanceof Error ? error.message : error || "Unknown error").slice(
    0,
    2000
  );

  await db
    .update(storageMigrationTask)
    .set({
      status: "failed",
      lastError: message,
      updatedAt: new Date(),
    })
    .where(eq(storageMigrationTask.id, id));
}

export async function getStorageTasksToCleanup({
  limit,
}: {
  limit: number;
}) {
  const now = new Date();
  return await db
    .select()
    .from(storageMigrationTask)
    .where(
      and(
        eq(storageMigrationTask.status, "migrated"),
        lte(storageMigrationTask.cleanupAfterAt, now)
      )
    )
    .orderBy(asc(storageMigrationTask.cleanupAfterAt))
    .limit(Math.max(1, Math.min(100, limit)));
}

export async function markStorageTaskCleaned(id: string) {
  await db
    .update(storageMigrationTask)
    .set({
      status: "deleted",
      cleanedAt: new Date(),
      lastError: null,
      updatedAt: new Date(),
    })
    .where(eq(storageMigrationTask.id, id));
}

export async function markStorageTaskCleanupFailed({
  id,
  error,
}: {
  id: string;
  error: unknown;
}) {
  const message = String(error instanceof Error ? error.message : error || "Unknown error").slice(
    0,
    2000
  );

  await db
    .update(storageMigrationTask)
    .set({
      lastError: message,
      updatedAt: new Date(),
    })
    .where(eq(storageMigrationTask.id, id));
}

export async function getStorageMigrationTaskStats() {
  const grouped = await db
    .select({
      status: storageMigrationTask.status,
      count: sql<number>`count(*)::int`,
    })
    .from(storageMigrationTask)
    .groupBy(storageMigrationTask.status);

  const pending = grouped
    .filter((item) => item.status === "pending" || item.status === "failed")
    .reduce((sum, item) => sum + (item.count || 0), 0);

  const cleanupPending = grouped
    .filter((item) => item.status === "migrated")
    .reduce((sum, item) => sum + (item.count || 0), 0);

  return {
    grouped,
    pending,
    cleanupPending,
  };
}