This commit is contained in:
Nicolas 2024-08-20 14:39:52 -03:00
parent c5ad4dedeb
commit b36faeaf54
3 changed files with 22 additions and 5 deletions

View File

@ -32,9 +32,14 @@ export async function crawlController(
await logCrawl(id, req.auth.team_id);
const { remainingCredits } = req.account;
// TODO: Get rid of crawlerOptions
const crawlerOptions = legacyCrawlerOptions(req.body.crawlerOptions);
const pageOptions = legacyScrapeOptions(req.body.scrapeOptions);
crawlerOptions.limit = Math.min(remainingCredits, crawlerOptions.limit);
const sc: StoredCrawl = {
originUrl: req.body.url,
crawlerOptions,

View File

@ -100,9 +100,10 @@ export type CrawlerOptions = z.infer<typeof crawlerOptions>;
export const crawlRequestSchema = z.object({
url,
origin: z.string().optional().default("api"),
crawlerOptions: crawlerOptions.default({}),
crawlerOptions: crawlerOptions.default({}), // TODO: Get rid of this
scrapeOptions: scrapeOptions.omit({ timeout: true }).default({}),
webhook: z.string().url().optional(),
limit: z.number().default(10000), //
});
// export type CrawlRequest = {
@ -225,20 +226,26 @@ type AuthObject = {
plan: string;
};
type Account = {
remainingCredits: number;
};
export interface RequestWithMaybeAuth<
ReqParams = {},
ReqBody = undefined,
ResBody = undefined
> extends Request<ReqParams, ReqBody, ResBody> {
auth?: AuthObject;
account?: Account;
}
export interface RequestWithAuth<
ReqParams = {},
ReqBody = undefined,
ResBody = undefined
ResBody = undefined,
> extends Request<ReqParams, ReqBody, ResBody> {
auth: AuthObject;
account?: Account;
}
export function legacyCrawlerOptions(x: CrawlerOptions) {

View File

@ -24,12 +24,17 @@ import { isUrlBlocked } from "../scraper/WebScraper/utils/blocklist";
// import { livenessController } from "../controllers/v1/liveness";
// import { readinessController } from "../controllers/v1/readiness";
function checkCreditsMiddleware(minimum: number): (req: RequestWithAuth, res: Response, next: NextFunction) => void {
function checkCreditsMiddleware(minimum?: number): (req: RequestWithAuth, res: Response, next: NextFunction) => void {
return (req, res, next) => {
(async () => {
if (!(await checkTeamCredits(req.auth.team_id, minimum)).success) {
if (!minimum && req.body) {
minimum = (req.body as any)?.limit ?? 1;
}
const { success, message, remainingCredits } = await checkTeamCredits(req.auth.team_id, minimum);
if (!success) {
return res.status(402).json({ success: false, error: "Insufficient credits" });
}
req.account = { remainingCredits }
next();
})()
.catch(err => next(err));
@ -101,7 +106,7 @@ v1Router.post(
blocklistMiddleware,
authMiddleware(RateLimiterMode.Crawl),
idempotencyMiddleware,
checkCreditsMiddleware(1),
checkCreditsMiddleware(),
wrap(crawlController)
);