Commit 24ac3314 authored by ThinhNC's avatar ThinhNC

feat(crawler): support real-time page persistence, partial export and prevent duplicate rerun

parent c917a34b
version: "3.8"
services:
postgres:
image: postgres:16-alpine
container_name: crawl_data_postgres
environment:
POSTGRES_USER: ${DB_USER:-postgres}
POSTGRES_PASSWORD: ${DB_PASSWORD:-postgres}
POSTGRES_DB: ${DB_NAME:-crawl_data_db}
ports:
- "${DB_PORT:-5432}:5432"
volumes:
- postgres_data:/var/lib/postgresql/data
# postgres:
# image: postgres:16-alpine
# container_name: crawl_data_postgres
# environment:
# POSTGRES_USER: ${DB_USER:-postgres}
# POSTGRES_PASSWORD: ${DB_PASSWORD:-postgres}
# POSTGRES_DB: ${DB_NAME:-crawl_data_db}
# ports:
# - "${DB_PORT:-5432}:5432"
# volumes:
# - postgres_data:/var/lib/postgresql/data
redis:
image: redis:7-alpine
......
......@@ -94,7 +94,7 @@ describe("CrawlExportService", () => {
expect(result).toEqual(exportRecord);
});
it("throws 400 when job is not COMPLETED", async () => {
it("throws 400 when job is still RUNNING", async () => {
(prisma.crawlJob.findUnique as jest.Mock).mockResolvedValue(
makeJob({ status: "RUNNING" }),
);
......@@ -104,6 +104,58 @@ describe("CrawlExportService", () => {
).rejects.toMatchObject({ statusCode: 400 });
});
it("allows exporting a CANCELED job when successPages > 0", async () => {
(prisma.crawlJob.findUnique as jest.Mock).mockResolvedValue(
makeJob({ status: "CANCELED", successPages: 5 }),
);
const exportRecord = makeExport();
mockExportService.generate.mockResolvedValue(exportRecord as any);
const result = await service.createExport(
"user-1",
"CRAWLER_USER",
"job-1",
"JSON",
);
expect(mockExportService.generate).toHaveBeenCalledWith(
expect.objectContaining({ id: "job-1" }),
"JSON",
);
expect(result).toEqual(exportRecord);
});
it("throws 400 when CANCELED job has 0 successPages", async () => {
(prisma.crawlJob.findUnique as jest.Mock).mockResolvedValue(
makeJob({ status: "CANCELED", successPages: 0 }),
);
await expect(
service.createExport("user-1", "CRAWLER_USER", "job-1", "JSON"),
).rejects.toMatchObject({ statusCode: 400 });
});
it("allows exporting a FAILED job when successPages > 0", async () => {
(prisma.crawlJob.findUnique as jest.Mock).mockResolvedValue(
makeJob({ status: "FAILED", successPages: 10 }),
);
const exportRecord = makeExport();
mockExportService.generate.mockResolvedValue(exportRecord as any);
const result = await service.createExport(
"user-1",
"CRAWLER_USER",
"job-1",
"CSV",
);
expect(mockExportService.generate).toHaveBeenCalledWith(
expect.objectContaining({ id: "job-1" }),
"CSV",
);
expect(result).toEqual(exportRecord);
});
it("throws 404 when job does not exist", async () => {
(prisma.crawlJob.findUnique as jest.Mock).mockResolvedValue(null);
......
......@@ -61,11 +61,28 @@ export class CrawlExportService {
);
}
if (job.status !== JOB_STATUS.COMPLETED) {
const isExportable =
job.status === JOB_STATUS.COMPLETED ||
(job.status === JOB_STATUS.CANCELED && (job.successPages ?? 0) > 0) ||
(job.status === JOB_STATUS.FAILED && (job.successPages ?? 0) > 0);
if (!isExportable) {
if (
job.status === JOB_STATUS.RUNNING ||
job.status === JOB_STATUS.PENDING ||
job.status === JOB_STATUS.QUEUED ||
job.status === JOB_STATUS.PROCESSING_EXPORT
) {
throw new AppError(
"Crawl job is still in progress",
400,
ERROR_CODE.CRAWL_JOB_NOT_COMPLETED,
);
}
throw new AppError(
"Crawl job is not completed yet",
"No successfully crawled pages available to export",
400,
ERROR_CODE.CRAWL_JOB_NOT_COMPLETED,
ERROR_CODE.VALIDATION_ERROR,
);
}
......
......@@ -14,7 +14,7 @@ describe("CrawlJobController - SSE Events", () => {
findById: jest.fn(),
} as any;
(CrawlJobService as jest.Mock).mockReturnValue(mockService);
(CrawlJobService as unknown as jest.Mock).mockReturnValue(mockService);
controller = new CrawlJobController();
});
......
......@@ -389,4 +389,19 @@ export class CrawlJobRepository {
]);
return { items, total, page: safePage, limit: safeLimit };
}
findRecentActiveJob(userId: string, startUrl: string, windowMs = 5000) {
const since = new Date(Date.now() - windowMs);
return prisma.crawlJob.findFirst({
where: {
userId,
startUrl,
status: {
in: [JOB_STATUS.PENDING, JOB_STATUS.QUEUED, JOB_STATUS.RUNNING],
},
createdAt: { gte: since },
},
orderBy: { createdAt: "desc" },
});
}
}
......@@ -165,7 +165,7 @@ export class CrawlJobService {
);
}
await crawlQueue.add("crawl-job", { jobId: job.id });
await crawlQueue.add("crawl-job", { jobId: job.id }, { jobId: job.id });
return job;
}
......@@ -216,6 +216,25 @@ export class CrawlJobService {
JOB_STATUS.CANCELED,
);
// Remove from BullMQ queue if still waiting/delayed
if (crawlQueue) {
try {
const bullJob = await crawlQueue.getJob(jobId);
if (bullJob) {
await bullJob.remove();
} else {
const waitingJobs = await crawlQueue.getJobs(["waiting", "delayed", "prioritized"]);
for (const wj of waitingJobs) {
if (wj.data?.jobId === jobId) {
await wj.remove();
}
}
}
} catch {
// Ignored
}
}
// For CRAWL mode: also cancel at the Firecrawl provider level to stop
// quota consumption. firecrawlJobId is saved by the worker as soon as
// asyncCrawlUrl() returns, so it may be null if the job was canceled
......@@ -288,20 +307,68 @@ export class CrawlJobService {
await storage.deleteFile(job.diffReportPath).catch(() => {});
}
if (crawlQueue) {
try {
const bullJob = await crawlQueue.getJob(jobId);
if (bullJob) {
await bullJob.remove();
} else {
const waitingJobs = await crawlQueue.getJobs(["waiting", "delayed", "prioritized"]);
for (const wj of waitingJobs) {
if (wj.data?.jobId === jobId) {
await wj.remove();
}
}
}
} catch {
// Ignored
}
}
await this.repository.delete(jobId);
return { success: true, message: "Crawl job deleted successfully" };
}
private static readonly rerunLocks = new Set<string>();
async rerun(userId: string, role: string, jobId: string) {
const existing = await this.findById(userId, role, jobId);
return this.create(userId, {
startUrl: existing.startUrl,
mode: existing.mode,
maxPages: existing.maxPages,
maxDepth: existing.maxDepth,
urls: existing.urls,
});
const lockKey = `${userId}:${jobId}`;
if (CrawlJobService.rerunLocks.has(lockKey)) {
if (this.repository.findRecentActiveJob) {
const recent = await this.repository.findRecentActiveJob(
userId,
existing.startUrl,
10000,
);
if (recent) return recent;
}
}
if (this.repository.findRecentActiveJob) {
const recent = await this.repository.findRecentActiveJob(
userId,
existing.startUrl,
5000,
);
if (recent) {
return recent;
}
}
CrawlJobService.rerunLocks.add(lockKey);
try {
return await this.create(userId, {
startUrl: existing.startUrl,
mode: existing.mode,
maxPages: existing.maxPages,
maxDepth: existing.maxDepth,
urls: existing.urls,
});
} finally {
setTimeout(() => CrawlJobService.rerunLocks.delete(lockKey), 3000);
}
}
async getLogs(
......
......@@ -111,7 +111,11 @@ export class FirecrawlService {
url: string,
maxPages: number,
maxDepth: number,
onProgress?: (completed: number, total: number) => void | Promise<void>,
onProgress?: (
completed: number,
total: number,
currentPages?: FirecrawlPageResult[],
) => void | Promise<void>,
shouldCancel?: () => Promise<boolean>,
): Promise<CrawlStatusResult> {
const client = getFirecrawlClient();
......@@ -171,7 +175,16 @@ export class FirecrawlService {
};
}
await onProgress?.(status.completed, status.total);
let currentPages: FirecrawlPageResult[] | undefined;
if (status.data && status.data.length > 0) {
currentPages = status.data.map((doc) => normalizePage(doc, url));
}
if (currentPages) {
await onProgress?.(status.completed, status.total, currentPages);
} else {
await onProgress?.(status.completed, status.total);
}
if (
status.status === "completed" ||
......@@ -292,7 +305,11 @@ export class FirecrawlService {
async batchScrapePages(
urls: string[],
maxPages: number,
onProgress?: (completed: number, total: number) => void | Promise<void>,
onProgress?: (
completed: number,
total: number,
currentPages?: FirecrawlPageResult[],
) => void | Promise<void>,
shouldCancel?: () => Promise<boolean>,
): Promise<CrawlStatusResult> {
const client = getFirecrawlClient();
......@@ -361,7 +378,18 @@ export class FirecrawlService {
};
}
await onProgress?.(status.completed, status.total);
let currentPages: FirecrawlPageResult[] | undefined;
if (status.data && status.data.length > 0) {
currentPages = status.data.map((doc) =>
normalizePage(doc as FirecrawlDocument, urls[0]),
);
}
if (currentPages) {
await onProgress?.(status.completed, status.total, currentPages);
} else {
await onProgress?.(status.completed, status.total);
}
if (
status.status === "completed" ||
......
......@@ -488,3 +488,95 @@ describe("processCrawlJob — URL_LIST mode", () => {
expect(mockPageRepo.upsert).toHaveBeenCalledTimes(1);
});
});
// ── Real-time persistence & Cancel preservation ────────────────────────────
describe("processCrawlJob — Real-time persistence & Cancel preservation", () => {
it("persists pages incrementally during live crawl onProgress calls", async () => {
mockJobRepo.findById = jest.fn().mockResolvedValue(makeJob({ mode: "CRAWL" }));
mockFirecrawl.crawlSite = jest
.fn()
.mockImplementation(async (_url, _maxPages, _maxDepth, onProgress) => {
// First progress event with 1 page
await onProgress?.(1, 2, [
{
url: "https://example.com/p1",
success: true,
markdown: "# P1",
metadata: { statusCode: 200 },
},
]);
return {
success: true,
status: "completed",
pages: [
{
url: "https://example.com/p1",
success: true,
markdown: "# P1",
metadata: { statusCode: 200 },
},
{
url: "https://example.com/p2",
success: true,
markdown: "# P2",
metadata: { statusCode: 200 },
},
],
failedUrls: [],
robotsBlockedUrls: [],
total: 2,
};
});
await processCrawlJob(makeBullJob("job-1"));
// Upsert called twice for 2 pages
expect(mockPageRepo.upsert).toHaveBeenCalledTimes(2);
expect(mockJobRepo.updateStatus).toHaveBeenCalledWith(
"job-1",
"COMPLETED",
expect.objectContaining({
successPages: 2,
totalPages: 2,
}),
);
});
it("preserves scraped pages and maintains CANCELED status when cancelled during crawl", async () => {
mockJobRepo.findById = jest.fn().mockResolvedValue(makeJob({ mode: "CRAWL" }));
mockFirecrawl.crawlSite = jest
.fn()
.mockImplementation(async (_url, _maxPages, _maxDepth, onProgress) => {
await onProgress?.(1, 10, [
{
url: "https://example.com/p1",
success: true,
markdown: "# P1",
metadata: { statusCode: 200 },
},
]);
return {
success: false,
status: "cancelled",
error: "Cancelled locally",
pages: [],
total: 10,
};
});
await processCrawlJob(makeBullJob("job-1"));
// Page p1 was persisted during onProgress
expect(mockPageRepo.upsert).toHaveBeenCalledTimes(1);
// Preserved CANCELED status instead of overwriting to FAILED
expect(mockJobRepo.updateStatus).toHaveBeenCalledWith(
"job-1",
"CANCELED",
expect.objectContaining({
successPages: 1,
}),
);
});
});
This diff is collapsed.
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment