Skip to main content

mocra_core/common/model/
response.rs

1use crate::cacheable::CacheAble;
2use crate::common::model::Cookies;
3use crate::common::model::ExecutionMark;
4use crate::common::model::meta::MetaData;
5use uuid::Uuid;
6
7use serde::{Deserialize, Serialize};
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct Response {
10    pub id: Uuid,
11    pub platform: String,
12    pub account: String,
13    pub module: String,
14    pub status_code: u16,
15    pub cookies: Cookies,
16    pub content: Vec<u8>,
17    pub storage_path: Option<String>,
18    pub headers: Vec<(String, String)>,
19    pub task_retry_times: usize,
20    pub metadata: MetaData,
21    pub download_middleware: Vec<String>,
22    pub data_middleware: Vec<String>,
23    pub task_finished: bool,
24    pub context: ExecutionMark,
25    pub run_id: Uuid,
26    pub prefix_request: Uuid,
27    pub request_hash: Option<String>,
28    #[serde(default)]
29    pub priority: crate::common::model::Priority,
30}
31impl Response {
32    /// Decodes the response body as a UTF-8 string (strict; invalid bytes return an error).
33    pub fn text(&self) -> crate::errors::Result<&str> {
34        std::str::from_utf8(&self.content).map_err(|e| {
35            crate::errors::Error::new(
36                crate::errors::ErrorKind::Response,
37                Some(format!("response body is not valid UTF-8: {e}")),
38            )
39        })
40    }
41
42    /// Decodes the response body as UTF-8 (invalid bytes are replaced with U+FFFD).
43    pub fn text_lossy(&self) -> std::borrow::Cow<'_, str> {
44        String::from_utf8_lossy(&self.content)
45    }
46
47    /// Deserializes the response body as JSON.
48    pub fn json<T: serde::de::DeserializeOwned>(&self) -> crate::errors::Result<T> {
49        serde_json::from_slice(&self.content).map_err(|e| {
50            crate::errors::Error::new(
51                crate::errors::ErrorKind::Response,
52                Some(format!("response body is not valid JSON: {e}")),
53            )
54        })
55    }
56
57    pub fn task_id(&self) -> String {
58        format!("{}-{}", self.account, self.platform)
59    }
60    pub fn module_id(&self) -> String {
61        format!("{}-{}-{}", self.account, self.platform, self.module)
62    }
63    /// Run-scoped task identifier for error tracking.
64    pub fn task_runtime_id(&self) -> String {
65        format!("{}:{}:{}", self.platform, self.account, self.run_id)
66    }
67    /// Run-scoped module identifier for error tracking.
68    pub fn module_runtime_id(&self) -> String {
69        format!(
70            "{}-{}-{}-{}",
71            self.account, self.platform, self.module, self.run_id
72        )
73    }
74    pub fn get_meta<T>(&self, key: &str) -> Option<T>
75    where
76        T: for<'de> Deserialize<'de>,
77    {
78        self.metadata.get_trait_config::<T>(key)
79    }
80    pub fn get_login_config<T>(&self, key: &str) -> Option<T>
81    where
82        T: for<'de> Deserialize<'de>,
83    {
84        self.metadata.get_login_config::<T>(key)
85    }
86    pub fn get_module_config<T>(&self, key: &str) -> Option<T>
87    where
88        T: for<'de> Deserialize<'de>,
89    {
90        self.metadata.get_module_config::<T>(key)
91    }
92    pub fn get_task_config<T>(&self, key: &str) -> Option<T>
93    where
94        T: for<'de> Deserialize<'de>,
95    {
96        self.metadata.get_task_config::<T>(key)
97    }
98    pub fn get_context(&self) -> &ExecutionMark {
99        &self.context
100    }
101    pub fn with_context(mut self, ctx: ExecutionMark) -> Self {
102        self.context = ctx;
103        self
104    }
105}
106
107impl CacheAble for Response {
108    fn field() -> impl AsRef<str> {
109        "response"
110    }
111
112    fn serialized_size_hint(&self) -> Option<usize> {
113        Some(self.content.len() + self.headers.len() * 64)
114    }
115
116    fn clone_for_serialize(&self) -> Option<Self> {
117        Some(self.clone())
118    }
119}
120
121impl crate::common::model::priority::Prioritizable for Response {
122    fn get_priority(&self) -> crate::common::model::Priority {
123        self.priority
124    }
125}
126
127use async_trait::async_trait;
128#[async_trait]
129impl crate::common::interface::storage::Offloadable for Response {
130    fn should_offload(&self, threshold: usize) -> bool {
131        self.content.len() > threshold && self.storage_path.is_none()
132    }
133
134    async fn offload(
135        &mut self,
136        storage: &std::sync::Arc<dyn crate::common::interface::storage::BlobStorage>,
137    ) -> crate::errors::Result<()> {
138        if self.content.is_empty() {
139            return Ok(());
140        }
141        let key = format!("response/{}/{}.bin", self.run_id, self.id);
142        let path = storage.put(&key, &self.content).await?;
143        self.storage_path = Some(path);
144        self.content.clear();
145        self.content.shrink_to_fit();
146        Ok(())
147    }
148
149    async fn reload(
150        &mut self,
151        storage: &std::sync::Arc<dyn crate::common::interface::storage::BlobStorage>,
152    ) -> crate::errors::Result<()> {
153        if let Some(path) = &self.storage_path {
154            // If content is already present (e.g. reloaded twice), skip?
155            // But if we want to ensure full content, we should check if empty.
156            if self.content.is_empty() {
157                let data = storage.get(path).await?;
158                self.content = data;
159            }
160        }
161        Ok(())
162    }
163}