spire_core/
context.rs

1//! The [`Context`] a worker handles.
2
3use std::fmt;
4use std::sync::Arc;
5
6use crate::Result;
7use crate::backend::Client;
8use crate::dataset::{Data, Dataset, Datasets};
9use crate::frontier::memory::InMemFrontier;
10use crate::frontier::queue::Queue;
11use crate::request::{Request, Response, ResponseExt};
12
13/// A request being crawled, with everything needed to handle it.
14///
15/// It holds the [`Request`], a client to fetch it with, a [`Queue`] to add
16/// the links it leads to, and the crawl's datasets.
17pub struct Context<C> {
18    request: Request,
19    client: C,
20    queue: Queue,
21    datasets: Arc<Datasets>,
22}
23
24impl<C> Context<C> {
25    /// Creates a standalone context, e.g. to test a handler.
26    ///
27    /// It gets an empty queue of its own and no datasets; see
28    /// [`with_queue`](Self::with_queue) and
29    /// [`with_dataset`](Self::with_dataset).
30    #[must_use]
31    pub fn new(request: Request, client: C) -> Self {
32        let queue = Queue::new(InMemFrontier::fifo()).child_of(&request);
33        Self::from_parts(request, client, queue, Arc::default())
34    }
35
36    pub(crate) const fn from_parts(
37        request: Request,
38        client: C,
39        queue: Queue,
40        datasets: Arc<Datasets>,
41    ) -> Self {
42        Self {
43            request,
44            client,
45            queue,
46            datasets,
47        }
48    }
49
50    /// Replaces the queue, binding it to the request.
51    #[must_use]
52    pub fn with_queue(mut self, queue: &Queue) -> Self {
53        self.queue = queue.child_of(&self.request);
54        self
55    }
56
57    /// Registers `dataset` for items of type `T`.
58    #[must_use]
59    pub fn with_dataset<T, D>(mut self, dataset: D) -> Self
60    where
61        T: Send + 'static,
62        D: Dataset<T>,
63    {
64        Arc::make_mut(&mut self.datasets).insert(dataset);
65        self
66    }
67
68    /// Returns the request.
69    #[must_use]
70    pub const fn request(&self) -> &Request {
71        &self.request
72    }
73
74    /// Returns the request mutably, e.g. to add headers before fetching.
75    pub const fn request_mut(&mut self) -> &mut Request {
76        &mut self.request
77    }
78
79    /// Returns the client.
80    #[must_use]
81    pub const fn client(&self) -> &C {
82        &self.client
83    }
84
85    /// Returns the client mutably.
86    pub const fn client_mut(&mut self) -> &mut C {
87        &mut self.client
88    }
89
90    /// Returns the queue for the links this request leads to.
91    #[must_use]
92    pub const fn queue(&self) -> &Queue {
93        &self.queue
94    }
95
96    /// Returns the dataset for items of type `T`.
97    ///
98    /// # Errors
99    ///
100    /// Returns a [`Dataset`](crate::ErrorKind::Dataset) error if no dataset
101    /// is registered for `T`.
102    pub fn dataset<T: Send + 'static>(&self) -> Result<Data<T>> {
103        self.datasets.get()
104    }
105
106    /// Splits the context into its request and client.
107    #[must_use]
108    pub fn into_parts(self) -> (Request, C) {
109        (self.request, self.client)
110    }
111
112    /// Fetches the response for the request.
113    ///
114    /// Links followed through the queue, including clones taken earlier,
115    /// resolve against the URI the response was served from from then on.
116    ///
117    /// # Errors
118    ///
119    /// Returns an error if the client fails to fetch the response.
120    pub async fn fetch(self) -> Result<Response>
121    where
122        C: Client,
123    {
124        let response = self.client.fetch(self.request).await?;
125        if let Some(uri) = response.final_uri() {
126            self.queue.rebase(uri.clone());
127        }
128        Ok(response)
129    }
130}
131
132impl<C> fmt::Debug for Context<C> {
133    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
134        f.debug_struct("Context")
135            .field("method", self.request.method())
136            .field("uri", self.request.uri())
137            .field("queue", &self.queue)
138            .finish_non_exhaustive()
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use crate::backend::Noop;
146    use crate::dataset::InMemDataset;
147    use crate::frontier::seed::Seed;
148    use crate::request::{Depth, RequestExt};
149
150    #[tokio::test]
151    async fn standalone_context() {
152        let request = Seed::get("https://example.com/").unwrap().into_request();
153        let cx = Context::new(request, Noop).with_dataset(InMemDataset::<u32>::new());
154
155        cx.dataset::<u32>().unwrap().write(1).await.unwrap();
156        assert!(cx.dataset::<String>().is_err());
157
158        cx.queue().follow("https://example.com/next").await.unwrap();
159        assert_eq!(cx.queue().size_hint(), Some(1));
160        assert_eq!(cx.request().depth(), Depth::ROOT);
161
162        let response = cx.fetch().await.unwrap();
163        assert!(response.status().is_success());
164    }
165}
spire_core/context.rs

spire_core/
context.rs
