1use 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#[derive(Debug, Clone, Serialize, Deserialize, Default)]
17pub struct TimestreamWriteStateView {
18 #[serde(default)]
19 pub databases: HashMap<String, DatabaseView>,
20 #[serde(default)]
22 pub tables: HashMap<String, TableView>,
23 #[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 pub creation_time: String,
36 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 pub creation_time: String,
53 pub last_updated_time: String,
55 #[serde(default)]
56 pub tags: HashMap<String, String>,
57 }
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 pub creation_time: String,
72 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
97impl 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
172impl 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
256impl 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}