106 lines
2.8 KiB
TypeScript
106 lines
2.8 KiB
TypeScript
import type { NextApiRequest, NextApiResponse } from 'next';
|
|
import { jsonRes } from '@/service/response';
|
|
import { connectToDatabase, User } from '@/service/mongo';
|
|
import { authUser } from '@/service/utils/auth';
|
|
import { PgDatasetTableName } from '@/constants/plugin';
|
|
import { findAllChildrenIds } from '../delete';
|
|
import QueryStream from 'pg-query-stream';
|
|
import Papa from 'papaparse';
|
|
|
|
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
|
|
try {
|
|
let { kbId } = req.query as {
|
|
kbId: string;
|
|
};
|
|
|
|
if (!kbId || !global.pgClient) {
|
|
throw new Error('缺少参数');
|
|
}
|
|
|
|
await connectToDatabase();
|
|
|
|
// 凭证校验
|
|
const { userId } = await authUser({ req, authToken: true });
|
|
|
|
const exportIds = [kbId, ...(await findAllChildrenIds(kbId))];
|
|
|
|
const thirtyMinutesAgo = new Date(
|
|
Date.now() - (global.feConfigs?.limit?.exportLimitMinutes || 0) * 60 * 1000
|
|
);
|
|
|
|
// auth export times
|
|
const authTimes = await User.findOne(
|
|
{
|
|
_id: userId,
|
|
$or: [
|
|
{ 'limit.exportKbTime': { $exists: false } },
|
|
{ 'limit.exportKbTime': { $lte: thirtyMinutesAgo } }
|
|
]
|
|
},
|
|
'_id limit'
|
|
);
|
|
|
|
if (!authTimes) {
|
|
const minutes = `${global.feConfigs?.limit?.exportLimitMinutes || 0} 分钟`;
|
|
throw new Error(`上次导出未到 ${minutes}，每 ${minutes}仅可导出一次。`);
|
|
}
|
|
|
|
// connect pg
|
|
global.pgClient.connect((err, client, done) => {
|
|
if (err) {
|
|
console.error(err);
|
|
res.end('Error connecting to database');
|
|
return;
|
|
}
|
|
// create pg select stream
|
|
const query = new QueryStream(
|
|
`SELECT q, a, source FROM ${PgDatasetTableName} where user_id='${userId}' AND kb_id IN (${exportIds
|
|
.map((id) => `'${id}'`)
|
|
.join(',')})`
|
|
);
|
|
const stream = client.query(query);
|
|
|
|
res.setHeader('Content-Disposition', 'attachment; filename=dataset.csv');
|
|
res.setHeader('Content-Type', 'text/csv');
|
|
|
|
res.write('index,content,source');
|
|
|
|
// parse data every row
|
|
stream.on('data', (row: { q: string; a: string; source?: string }) => {
|
|
const csv = Papa.unparse([row], { header: false });
|
|
res.write(`\n${csv}`);
|
|
});
|
|
stream.on('end', async () => {
|
|
try {
|
|
// update export time
|
|
await User.findByIdAndUpdate(userId, {
|
|
'limit.exportKbTime': new Date()
|
|
});
|
|
} catch (error) {}
|
|
|
|
// close response
|
|
done();
|
|
res.end();
|
|
});
|
|
stream.on('error', (err) => {
|
|
done(err);
|
|
res.end('Error exporting data');
|
|
});
|
|
});
|
|
} catch (err) {
|
|
res.status(500);
|
|
jsonRes(res, {
|
|
code: 500,
|
|
error: err
|
|
});
|
|
}
|
|
}
|
|
|
|
export const config = {
|
|
api: {
|
|
bodyParser: {
|
|
sizeLimit: '200mb'
|
|
}
|
|
}
|
|
};
|