mocra 0.3.0

A distributed, event-driven crawling and data collection framework
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

use crate::cacheable::CacheAble;
use crate::common::model::Cookies;
use crate::common::model::ExecutionMark;
use crate::common::model::meta::MetaData;
use uuid::Uuid;

use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Response {
    pub id: Uuid,
    pub platform: String,
    pub account: String,
    pub module: String,
    pub status_code: u16,
    pub cookies: Cookies,
    pub content: Vec<u8>,
    pub storage_path: Option<String>,
    pub headers: Vec<(String, String)>,
    pub task_retry_times: usize,
    pub metadata: MetaData,
    pub download_middleware: Vec<String>,
    pub data_middleware: Vec<String>,
    pub task_finished: bool,
    pub context: ExecutionMark,
    pub run_id: Uuid,
    pub prefix_request: Uuid,
    pub request_hash: Option<String>,
    #[serde(default)]
    pub priority: crate::common::model::Priority,
}
impl Response {
    pub fn task_id(&self) -> String {
        format!("{}-{}", self.account, self.platform)
    }
    pub fn module_id(&self) -> String {
        format!("{}-{}-{}", self.account, self.platform, self.module)
    }
    /// Run-scoped task identifier for error tracking.
    pub fn task_runtime_id(&self) -> String {
        format!("{}:{}:{}", self.platform, self.account, self.run_id)
    }
    /// Run-scoped module identifier for error tracking.
    pub fn module_runtime_id(&self) -> String {
        format!(
            "{}-{}-{}-{}",
            self.account, self.platform, self.module, self.run_id
        )
    }
    pub fn get_meta<T>(&self, key: &str) -> Option<T>
    where
        T: for<'de> Deserialize<'de>,
    {
        self.metadata.get_trait_config::<T>(key)
    }
    pub fn get_login_config<T>(&self, key: &str) -> Option<T>
    where
        T: for<'de> Deserialize<'de>,
    {
        self.metadata.get_login_config::<T>(key)
    }
    pub fn get_module_config<T>(&self, key: &str) -> Option<T>
    where
        T: for<'de> Deserialize<'de>,
    {
        self.metadata.get_module_config::<T>(key)
    }
    pub fn get_task_config<T>(&self, key: &str) -> Option<T>
    where
        T: for<'de> Deserialize<'de>,
    {
        self.metadata.get_task_config::<T>(key)
    }
    pub fn get_context(&self) -> &ExecutionMark {
        &self.context
    }
    pub fn with_context(mut self, ctx: ExecutionMark) -> Self {
        self.context = ctx;
        self
    }
}

impl CacheAble for Response {
    fn field() -> impl AsRef<str> {
        "response"
    }

    fn serialized_size_hint(&self) -> Option<usize> {
        Some(self.content.len() + self.headers.len() * 64)
    }

    fn clone_for_serialize(&self) -> Option<Self> {
        Some(self.clone())
    }
}

impl crate::common::model::priority::Prioritizable for Response {
    fn get_priority(&self) -> crate::common::model::Priority {
        self.priority
    }
}

use async_trait::async_trait;
#[async_trait]
impl crate::common::interface::storage::Offloadable for Response {
    fn should_offload(&self, threshold: usize) -> bool {
        self.content.len() > threshold && self.storage_path.is_none()
    }

    async fn offload(
        &mut self,
        storage: &std::sync::Arc<dyn crate::common::interface::storage::BlobStorage>,
    ) -> crate::errors::Result<()> {
        if self.content.is_empty() {
            return Ok(());
        }
        let key = format!("response/{}/{}.bin", self.run_id, self.id);
        let path = storage.put(&key, &self.content).await?;
        self.storage_path = Some(path);
        self.content.clear();
        self.content.shrink_to_fit();
        Ok(())
    }

    async fn reload(
        &mut self,
        storage: &std::sync::Arc<dyn crate::common::interface::storage::BlobStorage>,
    ) -> crate::errors::Result<()> {
        if let Some(path) = &self.storage_path {
            // If content is already present (e.g. reloaded twice), skip?
            // But if we want to ensure full content, we should check if empty.
            if self.content.is_empty() {
                let data = storage.get(path).await?;
                self.content = data;
            }
        }
        Ok(())
    }
}