mocra_core/common/model/
data.rs

1use crate::common::interface::StoreTrait;
2use crate::common::model::Response;
3use crate::common::model::meta::MetaData;
4#[cfg(feature = "polars")]
5use polars::io::SerWriter;
6#[cfg(feature = "polars")]
7use polars::io::ipc::IpcWriter;
8#[cfg(feature = "polars")]
9use polars::prelude::*;
10use std::fmt::Debug;
11use uuid::Uuid;
12/// Storage context carrying task/module identity.
13#[derive(Clone, Debug, Default)]
14pub struct StoreContext {
15    /// Unique request identifier.
16    pub request_id: Uuid,
17    /// Platform identifier.
18    pub platform: String,
19    /// Account identifier.
20    pub account: String,
21    /// Module identifier.
22    pub module: String,
23    /// Metadata.
24    pub meta: MetaData,
25    /// Data middleware list.
26    pub data_middleware: Vec<String>,
27}
28impl StoreContext {
29    /// Builds task identifier (`account-platform`).
30    pub fn task_id(&self) -> String {
31        format!("{}-{}", self.account, self.platform)
32    }
33    /// Builds module identifier (`account-platform-module`).
34    pub fn module_id(&self) -> String {
35        format!("{}-{}-{}", self.account, self.platform, self.module)
36    }
37}
38
39/// File storage structure for binary file content.
40#[derive(Clone, Debug, Default)]
41pub struct FileStore {
42    /// Context information.
43    pub ctx: StoreContext,
44    /// File name.
45    pub file_name: String,
46    /// File path.
47    pub file_path: String,
48    /// File content.
49    pub content: Vec<u8>,
50}
51
52impl StoreTrait for FileStore {
53    fn build(&self) -> DataEvent {
54        // Preserve all metadata by cloning the current FileStore into the enum variant.
55        DataEvent {
56            request_id: self.ctx.request_id,
57            platform: self.ctx.platform.clone(),
58            account: self.ctx.account.clone(),
59            module: self.ctx.module.clone(),
60            meta: self.ctx.meta.clone(),
61            data: DataType::File(self.clone()),
62            data_middleware: self.ctx.data_middleware.clone(),
63        }
64    }
65}
66
67impl From<FileStore> for DataEvent {
68    fn from(value: FileStore) -> Self {
69        let ctx_clone = value.ctx.clone();
70        DataEvent {
71            request_id: ctx_clone.request_id,
72            platform: ctx_clone.platform.clone(),
73            account: ctx_clone.account.clone(),
74            module: ctx_clone.module.clone(),
75            meta: ctx_clone.meta.clone(),
76            data: DataType::File(
77                FileStore::default()
78                    .with_ctx(ctx_clone.clone())
79                    .with_content(value.content)
80                    .with_name(value.file_name)
81                    .with_path(value.file_path),
82            ),
83            data_middleware: ctx_clone.data_middleware.clone(),
84        }
85    }
86}
87impl From<DataEvent> for FileStore {
88    fn from(value: DataEvent) -> Self {
89        match value.data {
90            DataType::File(f) => f,
91            _ => FileStore::default(),
92        }
93    }
94}
95
96impl FileStore {
97    /// Sets file content.
98    pub fn with_content(mut self, content: Vec<u8>) -> Self {
99        self.content = content;
100        self
101    }
102    /// Sets context information.
103    pub fn with_ctx(mut self, ctx: StoreContext) -> Self {
104        self.ctx = ctx;
105        self
106    }
107    /// Sets file name.
108    pub fn with_name(mut self, file_name: impl AsRef<str>) -> Self {
109        self.file_name = file_name.as_ref().to_string();
110        self
111    }
112    /// Sets file path.
113    pub fn with_path(mut self, file_path: impl AsRef<str>) -> Self {
114        self.file_path = file_path.as_ref().to_string();
115        self
116    }
117    /// Sets file name (alias).
118    pub fn with_file_name(self, file_name: impl AsRef<str>) -> Self {
119        self.with_name(file_name)
120    }
121    /// Sets file path (alias).
122    pub fn with_file_path(self, file_path: impl AsRef<str>) -> Self {
123        self.with_path(file_path)
124    }
125}
126
127#[cfg(feature = "polars")]
128#[derive(Clone, Debug)]
129pub enum DataframeStoreData {
130    Bytes(Vec<u8>),
131    DataFrame(DataFrame),
132}
133
134/// DataFrame storage structure for tabular data.
135#[cfg(feature = "polars")]
136#[derive(Clone, Debug)]
137pub struct DataFrameStore {
138    /// Context information.
139    ctx: StoreContext,
140    /// Serialized DataFrame payload (IPC format).
141    pub data: DataframeStoreData,
142    /// Database schema.
143    pub schema: String,
144    /// Database table name.
145    pub table: String,
146}
147#[cfg(feature = "polars")]
148impl Default for DataFrameStore {
149    fn default() -> Self {
150        Self {
151            ctx: StoreContext::default(),
152            data: DataframeStoreData::Bytes(vec![]),
153            schema: String::new(),
154            table: String::new(),
155        }
156    }
157}
158
159#[cfg(feature = "polars")]
160impl StoreTrait for DataFrameStore {
161    fn build(&self) -> DataEvent {
162        // Clone to preserve full metadata.
163        DataEvent {
164            request_id: self.ctx.request_id,
165            platform: self.ctx.platform.clone(),
166            account: self.ctx.account.clone(),
167            module: self.ctx.module.clone(),
168            meta: self.ctx.meta.clone(),
169            data: DataType::DataFrame(self.clone()),
170            data_middleware: self.ctx.data_middleware.clone(),
171        }
172    }
173}
174#[cfg(feature = "polars")]
175impl From<DataFrameStore> for DataEvent {
176    fn from(value: DataFrameStore) -> Self {
177        DataEvent {
178            request_id: value.ctx.request_id,
179            platform: value.ctx.platform.clone(),
180            account: value.ctx.account.clone(),
181            module: value.ctx.module.clone(),
182            meta: value.ctx.meta.clone(),
183            data_middleware: value.ctx.data_middleware.clone(),
184            data: DataType::DataFrame(value),
185        }
186    }
187}
188
189#[cfg(feature = "polars")]
190impl DataFrameStore {
191    /// Sets DataFrame data and serializes to IPC format.
192    pub fn with_data(mut self, data: DataFrame) -> Self {
193        let mut buffer = Vec::new();
194        let mut df = data;
195        let mut writer = IpcWriter::new(&mut buffer); // default options
196        writer.finish(&mut df).expect("serialize DataFrame to IPC");
197        self.data = DataframeStoreData::Bytes(buffer);
198        self
199    }
200    /// Sets schema.
201    pub fn with_schema(mut self, schema: impl AsRef<str>) -> Self {
202        self.schema = schema.as_ref().to_string();
203        self
204    }
205    /// Sets table name.
206    pub fn with_table(mut self, table: impl AsRef<str>) -> Self {
207        self.table = table.as_ref().to_string();
208        self
209    }
210    pub fn get_data(&self) -> Option<DataFrame> {
211        match &self.data {
212            DataframeStoreData::Bytes(bytes) => {
213                let cursor = std::io::Cursor::new(bytes);
214                let reader = polars::io::ipc::IpcReader::new(cursor);
215                reader.finish().ok()
216            }
217            DataframeStoreData::DataFrame(df) => Some(df.clone()),
218        }
219    }
220}
221
222/// Data type enum.
223// The variants are deliberately not boxed (`File`/`DataFrame` are inlined directly) so that
224// downstream code can destructure them straight away via `DataType::File(f)`; `Empty` is the
225// lightweight default. Boxing would change the public API, so the size-difference lint is
226// waived here.
227#[allow(clippy::large_enum_variant)]
228#[derive(Debug, Clone, Default)]
229pub enum DataType {
230    /// Empty payload: the default variant (no concrete data filled in yet), requiring no optional
231    /// dependencies.
232    #[default]
233    Empty,
234    /// Structured tabular data (requires the `polars` feature).
235    #[cfg(feature = "polars")]
236    DataFrame(DataFrameStore),
237    /// File data.
238    File(FileStore),
239}
240
241/// Generic data transfer object wrapping metadata and concrete payload.
242#[derive(Debug, Clone)]
243pub struct DataEvent {
244    /// Unique request identifier.
245    pub request_id: Uuid,
246    /// Platform identifier.
247    pub platform: String,
248    /// Account identifier.
249    pub account: String,
250    /// Module identifier.
251    pub module: String,
252    /// Metadata.
253    pub meta: MetaData,
254    /// Payload.
255    pub data: DataType,
256    /// Data middleware to execute.
257    pub data_middleware: Vec<String>,
258}
259
260impl Default for DataEvent {
261    fn default() -> Self {
262        Self {
263            request_id: Default::default(),
264            platform: "".to_string(),
265            account: "".to_string(),
266            module: "".to_string(),
267            meta: Default::default(),
268            data: DataType::Empty,
269            data_middleware: vec![],
270        }
271    }
272}
273impl DataEvent {
274    /// Builds `Data` from `Response`.
275    pub fn from(response: &Response) -> Self {
276        DataEvent {
277            request_id: response.id,
278            platform: response.platform.clone(),
279            account: response.account.clone(),
280            module: response.module.clone(),
281            meta: response.metadata.clone(),
282            data: DataType::Empty,
283            data_middleware: response.data_middleware.clone(),
284        }
285    }
286    /// Sets middleware list.
287    pub fn with_middlewares(mut self, middleware: Vec<String>) -> Self {
288        self.data_middleware = middleware;
289        self
290    }
291    /// Adds one middleware item.
292    pub fn with_middleware(mut self, middleware: impl AsRef<str>) -> Self {
293        self.data_middleware.push(middleware.as_ref().into());
294        self
295    }
296    /// Returns task ID (`account-platform`).
297    pub fn task_id(&self) -> String {
298        format!("{}-{}", self.account, self.platform)
299    }
300    /// Returns module ID (`account-platform-module`).
301    pub fn module_id(&self) -> String {
302        format!("{}-{}-{}", self.account, self.platform, self.module)
303    }
304    /// Converts into `DataFrameStore` (requires the `polars` feature).
305    #[cfg(feature = "polars")]
306    pub fn with_df(self, data: DataFrame) -> DataFrameStore {
307        DataFrameStore::default().with_data(data)
308    }
309    /// Converts into `FileStore`.
310    pub fn with_file(self, data: Vec<u8>) -> FileStore {
311        // Transfer CrawlData context into a FileStore builder preserving metadata.
312        FileStore {
313            ctx: StoreContext {
314                request_id: self.request_id,
315                platform: self.platform,
316                account: self.account,
317                module: self.module,
318                meta: self.meta,
319                data_middleware: self.data_middleware,
320            },
321            file_name: String::new(),
322            file_path: String::new(),
323            content: data,
324        }
325    }
326    /// Returns payload size (bytes or rows).
327    pub fn size(&self) -> usize {
328        match &self.data {
329            DataType::Empty => 0,
330            #[cfg(feature = "polars")]
331            DataType::DataFrame(df_store) => match &df_store.data {
332                DataframeStoreData::Bytes(bytes) => bytes.len(),
333                DataframeStoreData::DataFrame(df) => df.height() * df.width(), // rough estimate: rows * columns
334            },
335            DataType::File(file_store) => file_store.content.len(),
336        }
337    }
338}
339#[cfg(feature = "polars")]
340impl From<(DataFrame, &Response)> for DataEvent {
341    fn from(value: (DataFrame, &Response)) -> Self {
342        let (data, response) = value;
343        DataEvent {
344            request_id: response.id,
345            platform: response.platform.clone(),
346            account: response.account.clone(),
347            module: response.module.clone(),
348            meta: response.metadata.clone(),
349            data: DataType::DataFrame(DataFrameStore::default().with_data(data)),
350            data_middleware: response.data_middleware.clone(),
351        }
352    }
353}
354impl StoreTrait for DataEvent {
355    fn build(&self) -> DataEvent {
356        self.clone()
357    }
358}
359#[cfg(test)]
360mod tests {
361    use super::*;
362    use serde::{Deserialize, Serialize};
363
364    #[test]
365    fn test_file_store_builder() {
366        let store = FileStore::default()
367            .with_name("test.txt")
368            .with_path("/tmp/test.txt")
369            .with_content(vec![1, 2, 3]);
370
371        assert_eq!(store.file_name, "test.txt");
372        assert_eq!(store.file_path, "/tmp/test.txt");
373        assert_eq!(store.content, vec![1, 2, 3]);
374    }
375
376    #[test]
377    fn test_serde_examples() {
378        #[derive(Serialize, Deserialize, Debug)]
379        struct Item {
380            key: String,
381            value: serde_json::Value,
382        }
383
384        // Method 1: define only required fields; serde ignores unknown fields.
385        #[derive(Serialize, Deserialize, Debug)]
386        struct RespPartial {
387            data: Vec<Item>,
388            // `is_success` is not defined and will be ignored.
389        }
390
391        // Method 2: use `Option` for optional fields.
392        #[derive(Serialize, Deserialize, Debug)]
393        struct RespWithOptional {
394            data: Vec<Item>,
395            is_success: Option<bool>, // Optional field.
396            #[serde(default)] // Use default when field is missing.
397            extra_field: String,
398        }
399
400        // Test JSON with additional fields.
401        let complex_json = r#"
402        {
403            "data": [
404                {"key": "name", "value": "Alice"},
405                {"key": "age", "value": "30"},
406                {"key": "city", "value": "New York"}
407            ],
408            "is_success": true,
409            "timestamp": "2023-01-01T00:00:00Z",
410            "metadata": {
411                "total_count": 100,
412                "page": 1,
413                "limit": 10
414            },
415            "debug_info": "some debug data",
416            "version": "1.0.0"
417        }"#;
418
419        // Method 1: parse only required fields.
420        let resp_partial: RespPartial = serde_json::from_str(complex_json).unwrap();
421        assert_eq!(resp_partial.data.len(), 3);
422
423        // Method 2: handle optional fields via `Option`.
424        let resp_optional: RespWithOptional = serde_json::from_str(complex_json).unwrap();
425        assert_eq!(resp_optional.is_success, Some(true));
426        assert_eq!(resp_optional.extra_field, "");
427
428        // Method 4: dynamic parsing with `serde_json::Value`.
429        let value: serde_json::Value = serde_json::from_str(complex_json).unwrap();
430        if let Some(data) = value.get("data") {
431            let items: Vec<Item> = serde_json::from_value(data.clone()).unwrap();
432            assert_eq!(items.len(), 3);
433        }
434
435        // Extract specific fields.
436        let is_success = value
437            .get("is_success")
438            .and_then(|v| v.as_bool())
439            .unwrap_or(false);
440        let version = value
441            .get("version")
442            .and_then(|v| v.as_str())
443            .unwrap_or("unknown");
444
445        assert!(is_success);
446        assert_eq!(version, "1.0.0");
447    }
448
449    #[cfg(feature = "polars")]
450    #[test]
451    fn test_polars_dataframe() {
452        use chrono::NaiveDate;
453        let df: DataFrame = df!(
454            "name" => ["Alice Archer", "Ben Brown", "Chloe Cooper", "Daniel Donovan"],
455            "birthdate" => [
456                NaiveDate::from_ymd_opt(1997, 1, 10).unwrap(),
457                NaiveDate::from_ymd_opt(1985, 2, 15).unwrap(),
458                NaiveDate::from_ymd_opt(1983, 3, 22).unwrap(),
459                NaiveDate::from_ymd_opt(1981, 4, 30).unwrap(),
460            ],
461            "weight" => [57.9, 72.5, 53.6, 83.1],  // (kg)
462            "height" => [1.56, 1.77, 1.65, 1.75],  // (m)
463        )
464        .unwrap();
465
466        assert_eq!(df.height(), 4);
467        assert_eq!(df.width(), 4);
468    }
469}
mocra_core/common/model/data.rs

mocra_core/common/model/
data.rs
