@ -2,12 +2,38 @@ import fs from "node:fs/promises";
import path from "node:path" ;
import { dbGlobal } from "drizzle-pkg/lib/db" ;
import { userExportTasks } from "drizzle-pkg/lib/schema/export" ;
import { and , desc , eq , or } from "drizzle-orm" ;
import { and , desc , eq , or , sql } from "drizzle-orm" ;
import { nextIntegerId } from "../../utils/sqlite-id" ;
import { RELATIVE_TMP_DIR } from "#server/constants/media" ;
type ExportMaskPolicy = "masked" | "raw" ;
function positiveIntFromEnv ( name : string , fallback : number ) : number {
const raw = process . env [ name ] ;
if ( ! raw ) {
return fallback ;
}
const n = Number ( raw ) ;
return Number . isInteger ( n ) && n > 0 ? n : fallback ;
}
function positiveBigIntFromEnv ( name : string , fallback : bigint ) : bigint {
const raw = process . env [ name ] ;
if ( ! raw ) {
return fallback ;
}
try {
const n = BigInt ( raw ) ;
return n > 0 n ? n : fallback ;
} catch {
return fallback ;
}
}
const EXPORT_MAX_RUNNING_TASKS = positiveIntFromEnv ( "EXPORT_MAX_RUNNING_TASKS" , 2 ) ;
const EXPORT_MAX_QUEUED_TASKS = positiveIntFromEnv ( "EXPORT_MAX_QUEUED_TASKS" , 30 ) ;
const EXPORT_MAX_RETAINED_BYTES = positiveBigIntFromEnv ( "EXPORT_MAX_RETAINED_BYTES" , 2 n * 1024 n * 1024 n * 1024 n ) ;
function exportRootDir ( ) : string {
return path . resolve ( process . cwd ( ) , RELATIVE_TMP_DIR , "exports" ) ;
}
@ -61,6 +87,52 @@ export async function createExportTask(params: { userId: number; maskPolicy: Exp
if ( activeTask ) {
throw { statusCode : 409 , statusMessage : "已有导出任务在处理中,请稍后再试" } ;
}
const runningRows = await dbGlobal
. select ( { runningCount : sql < number > ` count(*) ` } )
. from ( userExportTasks )
. where ( eq ( userExportTasks . status , "running" ) ) ;
const runningCount = Number ( runningRows [ 0 ] ? . runningCount ? ? 0 ) ;
if ( runningCount >= EXPORT_MAX_RUNNING_TASKS ) {
throw {
statusCode : 429 ,
statusMessage : "导出任务繁忙,请稍后重试" ,
} ;
}
const queuedRows = await dbGlobal
. select ( { queuedCount : sql < number > ` count(*) ` } )
. from ( userExportTasks )
. where ( eq ( userExportTasks . status , "queued" ) ) ;
const queuedCount = Number ( queuedRows [ 0 ] ? . queuedCount ? ? 0 ) ;
if ( queuedCount >= EXPORT_MAX_QUEUED_TASKS ) {
throw {
statusCode : 429 ,
statusMessage : "导出排队过多,请稍后重试" ,
} ;
}
const now = new Date ( ) ;
const retainedRows = await dbGlobal
. select ( {
retainedBytes : sql < string > ` coalesce(sum( ${ userExportTasks . totalBytes } ), 0) ` ,
} )
. from ( userExportTasks )
. where (
and (
eq ( userExportTasks . status , "succeeded" ) ,
sql ` ${ userExportTasks . expiresAt } IS NOT NULL ` ,
sql ` ${ userExportTasks . expiresAt } > ${ now } ` ,
) ,
) ;
const retainedBytes = retainedRows [ 0 ] ? . retainedBytes ? ? "0" ;
if ( BigInt ( retainedBytes ) >= EXPORT_MAX_RETAINED_BYTES ) {
throw {
statusCode : 507 ,
statusMessage : "导出空间已达上限,请先删除旧导出后重试" ,
} ;
}
const id = await nextIntegerId ( userExportTasks , userExportTasks . id ) ;
await dbGlobal . insert ( userExportTasks ) . values ( {
id ,