pagesift_engine/engine/
hooked.rs

1//! 单页出口套上钩子（`pagesift_hooks`）：抓取前改请求或跳过、抓取后重试一次或丢弃、抽取后改文档。
2//! 没有钩子时直接走原路径（一次 `is_empty` 判断）。
3
4use pagesift_api::types::{Document, ScrapeOptions};
5use pagesift_api::{ApiError, ErrorCode};
6use pagesift_hooks::{Request, Verdict};
7
8use super::{Scraper, Source};
9
10impl Scraper {
11    pub(super) async fn scrape_inner(
12        &self,
13        source: Source<'_>,
14        opts: &ScrapeOptions,
15    ) -> Result<Document, ApiError> {
16        if self.hooks.is_empty() {
17            return self.scrape_plain(source, opts).await;
18        }
19        let (url, mode) = match source {
20            Source::Url(url, mode) => (url.to_string(), mode),
21            Source::Upload(_) => {
22                let doc = Box::pin(self.scrape_plain(source, opts)).await?;
23                let name = doc.metadata.source_url.clone().unwrap_or_default();
24                return Ok(self.hooks.run_document(&name, doc));
25            }
26        };
27        let req = self
28            .hooks
29            .run_request(Request {
30                url,
31                options: opts.clone(),
32            })
33            .map_err(|why| {
34                ApiError::new(
35                    ErrorCode::UnsupportedOption,
36                    format!("skipped by a request hook: {why}"),
37                )
38            })?;
39        let mut doc =
40            Box::pin(self.scrape_plain(Source::Url(&req.url, mode), &req.options)).await?;
41        match self.hooks.run_response(&req.url, &doc) {
42            Verdict::Keep => {}
43            // 只重试一次：钩子再要求重试也照用第二次的结果
44            Verdict::Retry => {
45                doc =
46                    Box::pin(self.scrape_plain(Source::Url(&req.url, mode), &req.options)).await?;
47            }
48            Verdict::Drop(why) => {
49                return Err(ApiError::new(
50                    ErrorCode::UpstreamError,
51                    format!("dropped by a response hook: {why}"),
52                ));
53            }
54        }
55        Ok(self.hooks.run_document(&req.url, doc))
56    }
57}
58
59/// `Scraper::builder(config).hook(...)`：在代码里注册钩子（和 `[[rules]]` 编译出的放在同一组链里）。
60pub struct ScraperBuilder {
61    config: super::ScraperConfig,
62}
63
64impl Scraper {
65    pub fn builder(config: super::ScraperConfig) -> ScraperBuilder {
66        ScraperBuilder { config }
67    }
68
69    /// 输出阶段的钩子由写结果的一方调用：Rust 用户在自己写结果的地方调；内置入口里只有 CLI 的 crawl 调用。
70    pub fn hooks(&self) -> &pagesift_hooks::Hooks {
71        &self.hooks
72    }
73}
74
75impl ScraperBuilder {
76    /// 注册一组钩子（`pagesift_hooks::Hooks` 上用 `request`、`response`、`document`、`sink` 加好）。
77    pub fn hook(mut self, hooks: &pagesift_hooks::Hooks) -> Self {
78        self.config.hooks.extend(hooks);
79        self
80    }
81
82    pub fn build(self) -> Result<Scraper, super::ScraperError> {
83        Scraper::new(self.config)
84    }
85}
pagesift_engine/engine/hooked.rs

pagesift_engine/engine/
hooked.rs
