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); await logCrawl(id, req.auth.team_id);
const { remainingCredits } = req.account;
// TODO: Get rid of crawlerOptions
const crawlerOptions = legacyCrawlerOptions(req.body.crawlerOptions); const crawlerOptions = legacyCrawlerOptions(req.body.crawlerOptions);
const pageOptions = legacyScrapeOptions(req.body.scrapeOptions); const pageOptions = legacyScrapeOptions(req.body.scrapeOptions);
crawlerOptions.limit = Math.min(remainingCredits, crawlerOptions.limit);
const sc: StoredCrawl = { const sc: StoredCrawl = {
originUrl: req.body.url, originUrl: req.body.url,
crawlerOptions, crawlerOptions,

View File

@ -100,9 +100,10 @@ export type CrawlerOptions = z.infer<typeof crawlerOptions>;
export const crawlRequestSchema = z.object({ export const crawlRequestSchema = z.object({
url, url,
origin: z.string().optional().default("api"), origin: z.string().optional().default("api"),
crawlerOptions: crawlerOptions.default({}), crawlerOptions: crawlerOptions.default({}), // TODO: Get rid of this
scrapeOptions: scrapeOptions.omit({ timeout: true }).default({}), scrapeOptions: scrapeOptions.omit({ timeout: true }).default({}),
webhook: z.string().url().optional(), webhook: z.string().url().optional(),
limit: z.number().default(10000), //
}); });
// export type CrawlRequest = { // export type CrawlRequest = {
@ -225,20 +226,26 @@ type AuthObject = {
plan: string; plan: string;
}; };
type Account = {
remainingCredits: number;
};
export interface RequestWithMaybeAuth< export interface RequestWithMaybeAuth<
ReqParams = {}, ReqParams = {},
ReqBody = undefined, ReqBody = undefined,
ResBody = undefined ResBody = undefined
> extends Request<ReqParams, ReqBody, ResBody> { > extends Request<ReqParams, ReqBody, ResBody> {
auth?: AuthObject; auth?: AuthObject;
account?: Account;
} }
export interface RequestWithAuth< export interface RequestWithAuth<
ReqParams = {}, ReqParams = {},
ReqBody = undefined, ReqBody = undefined,
ResBody = undefined ResBody = undefined,
> extends Request<ReqParams, ReqBody, ResBody> { > extends Request<ReqParams, ReqBody, ResBody> {
auth: AuthObject; auth: AuthObject;
account?: Account;
} }
export function legacyCrawlerOptions(x: CrawlerOptions) { 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 { livenessController } from "../controllers/v1/liveness";
// import { readinessController } from "../controllers/v1/readiness"; // 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) => { return (req, res, next) => {
(async () => { (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" }); return res.status(402).json({ success: false, error: "Insufficient credits" });
} }
req.account = { remainingCredits }
next(); next();
})() })()
.catch(err => next(err)); .catch(err => next(err));
@ -101,7 +106,7 @@ v1Router.post(
blocklistMiddleware, blocklistMiddleware,
authMiddleware(RateLimiterMode.Crawl), authMiddleware(RateLimiterMode.Crawl),
idempotencyMiddleware, idempotencyMiddleware,
checkCreditsMiddleware(1), checkCreditsMiddleware(),
wrap(crawlController) wrap(crawlController)
); );