Skip to main content

winterbaume_timestreamwrite/
views.rs

1//! Serde-compatible view types for Timestream Write state snapshots.
2
3use std::collections::HashMap;
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use winterbaume_core::{StateChangeNotifier, StateViewError, StatefulService};
8
9use crate::handlers::TimestreamWriteService;
10use crate::state::TimestreamWriteState;
11use crate::types::{
12    BatchLoadTask, Database, MagneticStoreWriteProperties, RetentionProperties, Table,
13};
14
15/// Serializable view of the entire Timestream Write state for one account/region.
16#[derive(Debug, Clone, Serialize, Deserialize, Default)]
17pub struct TimestreamWriteStateView {
18    #[serde(default)]
19    pub databases: HashMap<String, DatabaseView>,
20    /// Tables keyed by "database_name\x1ftable_name".
21    #[serde(default)]
22    pub tables: HashMap<String, TableView>,
23    /// Batch load tasks keyed by task_id.
24    #[serde(default)]
25    pub batch_load_tasks: HashMap<String, BatchLoadTaskView>,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct DatabaseView {
30    pub database_name: String,
31    pub arn: String,
32    pub kms_key_id: Option<String>,
33    pub table_count: i64,
34    /// Creation time in RFC 3339 format.
35    pub creation_time: String,
36    /// Last updated time in RFC 3339 format.
37    pub last_updated_time: String,
38    #[serde(default)]
39    pub tags: HashMap<String, String>,
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct TableView {
44    pub table_name: String,
45    pub database_name: String,
46    pub arn: String,
47    pub table_status: String,
48    pub memory_store_retention_period_in_hours: i64,
49    pub magnetic_store_retention_period_in_days: i64,
50    pub enable_magnetic_store_writes: bool,
51    /// Creation time in RFC 3339 format.
52    pub creation_time: String,
53    /// Last updated time in RFC 3339 format.
54    pub last_updated_time: String,
55    #[serde(default)]
56    pub tags: HashMap<String, String>,
57    // records are transient and not persisted
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize)]
61pub struct BatchLoadTaskView {
62    pub task_id: String,
63    pub task_status: String,
64    pub database_name: String,
65    pub table_name: String,
66    pub data_source_configuration: Value,
67    pub report_configuration: Value,
68    #[serde(default, skip_serializing_if = "Option::is_none")]
69    pub data_model_configuration: Option<Value>,
70    /// Creation time in RFC 3339 format.
71    pub creation_time: String,
72    /// Last updated time in RFC 3339 format.
73    pub last_updated_time: String,
74    #[serde(default, skip_serializing_if = "Option::is_none")]
75    pub error_message: Option<String>,
76}
77
78fn tuple2_key(a: &str, b: &str) -> String {
79    format!("{}\x1f{}", a, b)
80}
81
82fn parse_tuple2(key: &str) -> Option<(String, String)> {
83    let parts: Vec<&str> = key.splitn(2, '\x1f').collect();
84    if parts.len() == 2 {
85        Some((parts[0].to_string(), parts[1].to_string()))
86    } else {
87        None
88    }
89}
90
91fn parse_datetime(s: &str) -> chrono::DateTime<chrono::Utc> {
92    chrono::DateTime::parse_from_rfc3339(s)
93        .map(|dt| dt.with_timezone(&chrono::Utc))
94        .unwrap_or_else(|_| chrono::Utc::now())
95}
96
97// --- From internal types to view types ---
98
99impl From<&TimestreamWriteState> for TimestreamWriteStateView {
100    fn from(state: &TimestreamWriteState) -> Self {
101        TimestreamWriteStateView {
102            databases: state
103                .databases
104                .iter()
105                .map(|(k, v)| {
106                    (
107                        k.clone(),
108                        DatabaseView {
109                            database_name: v.database_name.clone(),
110                            arn: v.arn.clone(),
111                            kms_key_id: v.kms_key_id.clone(),
112                            table_count: v.table_count,
113                            creation_time: v.creation_time.to_rfc3339(),
114                            last_updated_time: v.last_updated_time.to_rfc3339(),
115                            tags: v.tags.clone(),
116                        },
117                    )
118                })
119                .collect(),
120            tables: state
121                .tables
122                .iter()
123                .map(|((db_name, tbl_name), v)| {
124                    (
125                        tuple2_key(db_name, tbl_name),
126                        TableView {
127                            table_name: v.table_name.clone(),
128                            database_name: v.database_name.clone(),
129                            arn: v.arn.clone(),
130                            table_status: v.table_status.clone(),
131                            memory_store_retention_period_in_hours: v
132                                .retention_properties
133                                .memory_store_retention_period_in_hours,
134                            magnetic_store_retention_period_in_days: v
135                                .retention_properties
136                                .magnetic_store_retention_period_in_days,
137                            enable_magnetic_store_writes: v
138                                .magnetic_store_write_properties
139                                .enable_magnetic_store_writes,
140                            creation_time: v.creation_time.to_rfc3339(),
141                            last_updated_time: v.last_updated_time.to_rfc3339(),
142                            tags: v.tags.clone(),
143                        },
144                    )
145                })
146                .collect(),
147            batch_load_tasks: state
148                .batch_load_tasks
149                .iter()
150                .map(|(k, v)| {
151                    (
152                        k.clone(),
153                        BatchLoadTaskView {
154                            task_id: v.task_id.clone(),
155                            task_status: v.task_status.clone(),
156                            database_name: v.database_name.clone(),
157                            table_name: v.table_name.clone(),
158                            data_source_configuration: v.data_source_configuration.clone(),
159                            report_configuration: v.report_configuration.clone(),
160                            data_model_configuration: v.data_model_configuration.clone(),
161                            creation_time: v.creation_time.to_rfc3339(),
162                            last_updated_time: v.last_updated_time.to_rfc3339(),
163                            error_message: v.error_message.clone(),
164                        },
165                    )
166                })
167                .collect(),
168        }
169    }
170}
171
172// --- From view types to internal types ---
173
174impl From<TimestreamWriteStateView> for TimestreamWriteState {
175    fn from(view: TimestreamWriteStateView) -> Self {
176        let databases = view
177            .databases
178            .into_iter()
179            .map(|(k, v)| {
180                (
181                    k,
182                    Database {
183                        database_name: v.database_name,
184                        arn: v.arn,
185                        kms_key_id: v.kms_key_id,
186                        table_count: v.table_count,
187                        creation_time: parse_datetime(&v.creation_time),
188                        last_updated_time: parse_datetime(&v.last_updated_time),
189                        tags: v.tags,
190                    },
191                )
192            })
193            .collect();
194
195        let tables = view
196            .tables
197            .into_iter()
198            .filter_map(|(k, v)| {
199                parse_tuple2(&k).map(|(db_name, tbl_name)| {
200                    (
201                        (db_name, tbl_name),
202                        Table {
203                            table_name: v.table_name,
204                            database_name: v.database_name,
205                            arn: v.arn,
206                            table_status: v.table_status,
207                            retention_properties: RetentionProperties {
208                                memory_store_retention_period_in_hours: v
209                                    .memory_store_retention_period_in_hours,
210                                magnetic_store_retention_period_in_days: v
211                                    .magnetic_store_retention_period_in_days,
212                            },
213                            magnetic_store_write_properties: MagneticStoreWriteProperties {
214                                enable_magnetic_store_writes: v.enable_magnetic_store_writes,
215                            },
216                            creation_time: parse_datetime(&v.creation_time),
217                            last_updated_time: parse_datetime(&v.last_updated_time),
218                            tags: v.tags,
219                            records: Vec::new(),
220                        },
221                    )
222                })
223            })
224            .collect();
225
226        let batch_load_tasks = view
227            .batch_load_tasks
228            .into_iter()
229            .map(|(k, v)| {
230                (
231                    k,
232                    BatchLoadTask {
233                        task_id: v.task_id,
234                        task_status: v.task_status,
235                        database_name: v.database_name,
236                        table_name: v.table_name,
237                        data_source_configuration: v.data_source_configuration,
238                        report_configuration: v.report_configuration,
239                        data_model_configuration: v.data_model_configuration,
240                        creation_time: parse_datetime(&v.creation_time),
241                        last_updated_time: parse_datetime(&v.last_updated_time),
242                        error_message: v.error_message,
243                    },
244                )
245            })
246            .collect();
247
248        TimestreamWriteState {
249            databases,
250            tables,
251            batch_load_tasks,
252        }
253    }
254}
255
256// --- StatefulService implementation ---
257
258impl StatefulService for TimestreamWriteService {
259    type StateView = TimestreamWriteStateView;
260
261    async fn snapshot(&self, account_id: &str, region: &str) -> Self::StateView {
262        let state = self.state.get(account_id, region);
263        let guard = state.read().await;
264        TimestreamWriteStateView::from(&*guard)
265    }
266
267    async fn restore(
268        &self,
269        account_id: &str,
270        region: &str,
271        view: Self::StateView,
272    ) -> Result<(), StateViewError> {
273        let state = self.state.get(account_id, region);
274        {
275            let mut guard = state.write().await;
276            *guard = TimestreamWriteState::from(view);
277        }
278        self.notify_state_changed(account_id, region).await;
279        Ok(())
280    }
281
282    async fn merge(
283        &self,
284        account_id: &str,
285        region: &str,
286        view: Self::StateView,
287    ) -> Result<(), StateViewError> {
288        let state = self.state.get(account_id, region);
289        {
290            let mut guard = state.write().await;
291            let new_state = TimestreamWriteState::from(view);
292            guard.databases.extend(new_state.databases);
293            guard.tables.extend(new_state.tables);
294            guard.batch_load_tasks.extend(new_state.batch_load_tasks);
295        }
296        self.notify_state_changed(account_id, region).await;
297        Ok(())
298    }
299
300    fn notifier(&self) -> &StateChangeNotifier<Self::StateView> {
301        &self.notifier
302    }
303}