parse_book_source/http_client/
mod.rs

1use crate::{HttpConfig, Result};
2use anyhow::anyhow;
3use rate_limiter::TokenBucket;
4use reqwest::{
5    Body, Client, ClientBuilder,
6    header::{HeaderMap, HeaderName, HeaderValue},
7};
8use std::time::Duration;
9pub mod rate_limiter;
10
11#[derive(Debug, Clone)]
12pub struct HttpClient {
13    pub client: Client,
14    pub base_url: String,
15    pub rate_limiter: Option<TokenBucket>,
16}
17
18impl HttpClient {
19    pub fn new(base_url: &str, config: &HttpConfig) -> Result<Self> {
20        let mut client = ClientBuilder::new().cookie_store(true);
21
22        if let Some(header) = &config.header {
23            let mut headers = HeaderMap::new();
24
25            for (k, v) in header {
26                headers.insert(
27                    HeaderName::try_from(k)
28                        .map_err(|e| anyhow!("header name is not valid: {}", e))?,
29                    HeaderValue::from_str(v)
30                        .map_err(|e| anyhow!("header value is not valid: {}", e))?,
31                );
32            }
33            client = client.default_headers(headers);
34        }
35
36        if let Some(timeout) = config.timeout {
37            client = client.timeout(Duration::from_millis(timeout));
38        }
39
40        Ok(Self {
41            client: client.build()?,
42            base_url: base_url.to_string(),
43            rate_limiter: config.rate_limit.as_ref().map(|rate_limit| {
44                TokenBucket::new(
45                    rate_limit.max_count as usize,
46                    Duration::from_secs_f64(rate_limit.fill_duration),
47                )
48            }),
49        })
50    }
51
52    fn url_with_base(&self, url: &str) -> String {
53        if url.starts_with("http") {
54            url.to_string()
55        } else if url.starts_with('/') {
56            format!("{}{}", self.base_url, url)
57        } else {
58            format!("{}/{}", self.base_url, url)
59        }
60    }
61
62    pub async fn get(&self, url: &str) -> Result<reqwest::Response> {
63        if let Some(rate_limiter) = &self.rate_limiter {
64            rate_limiter.acquire().await;
65        }
66
67        let url = self.url_with_base(url);
68
69        Ok(self.client.get(url).send().await?)
70    }
71
72    pub async fn post<T: Into<Body>>(&self, url: &str, body: T) -> Result<reqwest::Response> {
73        if let Some(rate_limiter) = &self.rate_limiter {
74            rate_limiter.acquire().await;
75        }
76
77        let url = self.url_with_base(url);
78
79        Ok(self.client.post(url).body(body).send().await?)
80    }
81}
parse_book_source/http_client/mod.rs

parse_book_source/http_client/
mod.rs
