winterbaume_timestreamwrite/
state.rs

1use std::collections::HashMap;
2
3use chrono::Utc;
4use serde_json::Value;
5use thiserror::Error;
6
7use crate::types::{
8    BatchLoadTask, Database, MagneticStoreWriteProperties, Record, RetentionProperties, Table,
9};
10
11/// In-memory state for the Timestream Write service.
12#[derive(Debug, Default)]
13pub struct TimestreamWriteState {
14    /// Databases keyed by database name.
15    pub databases: HashMap<String, Database>,
16    /// Tables keyed by (database_name, table_name).
17    pub tables: HashMap<(String, String), Table>,
18    /// Batch load tasks keyed by task_id.
19    pub batch_load_tasks: HashMap<String, BatchLoadTask>,
20}
21
22/// Error type for Timestream Write operations.
23#[derive(Debug, Error)]
24pub enum TimestreamWriteError {
25    #[error("Database {database_name} already exists in account {account_id} and region {region}")]
26    DatabaseAlreadyExists {
27        database_name: String,
28        account_id: String,
29        region: String,
30    },
31
32    #[error("Table {table_name} already exists in database {database_name}")]
33    TableAlreadyExists {
34        database_name: String,
35        table_name: String,
36    },
37
38    #[error(
39        "Database {database_name} contains tables. Delete all tables before deleting the database."
40    )]
41    DatabaseNotEmpty { database_name: String },
42
43    #[error(
44        "Task {task_id} cannot be resumed because its status is {task_status}. Only FAILED tasks can be resumed."
45    )]
46    TaskCannotBeResumed {
47        task_id: String,
48        task_status: String,
49    },
50
51    #[error("The database {database_name} does not exist.")]
52    DatabaseNotFound { database_name: String },
53
54    #[error("The table {table_name} does not exist in database {database_name}.")]
55    TableNotFound {
56        database_name: String,
57        table_name: String,
58    },
59
60    #[error("The resource {resource_arn} does not exist.")]
61    ResourceNotFound { resource_arn: String },
62
63    #[error("The task {task_id} does not exist.")]
64    TaskNotFound { task_id: String },
65}
66
67impl TimestreamWriteState {
68    // --- Database operations ---
69
70    pub fn create_database(
71        &mut self,
72        database_name: &str,
73        account_id: &str,
74        region: &str,
75        kms_key_id: Option<&str>,
76    ) -> Result<&Database, TimestreamWriteError> {
77        if self.databases.contains_key(database_name) {
78            return Err(TimestreamWriteError::DatabaseAlreadyExists {
79                database_name: database_name.to_string(),
80                account_id: account_id.to_string(),
81                region: region.to_string(),
82            });
83        }
84
85        let arn = format!("arn:aws:timestream:{region}:{account_id}:database/{database_name}");
86
87        let now = Utc::now();
88        let database = Database {
89            database_name: database_name.to_string(),
90            arn,
91            kms_key_id: kms_key_id.map(|s| s.to_string()),
92            table_count: 0,
93            creation_time: now,
94            last_updated_time: now,
95            tags: HashMap::new(),
96        };
97
98        self.databases.insert(database_name.to_string(), database);
99        Ok(self.databases.get(database_name).unwrap())
100    }
101
102    pub fn describe_database(
103        &self,
104        database_name: &str,
105    ) -> Result<&Database, TimestreamWriteError> {
106        self.databases
107            .get(database_name)
108            .ok_or_else(|| TimestreamWriteError::DatabaseNotFound {
109                database_name: database_name.to_string(),
110            })
111    }
112
113    pub fn delete_database(&mut self, database_name: &str) -> Result<(), TimestreamWriteError> {
114        if !self.databases.contains_key(database_name) {
115            return Err(TimestreamWriteError::DatabaseNotFound {
116                database_name: database_name.to_string(),
117            });
118        }
119        // Check if there are tables in the database
120        let has_tables = self.tables.keys().any(|(db, _)| db == database_name);
121        if has_tables {
122            return Err(TimestreamWriteError::DatabaseNotEmpty {
123                database_name: database_name.to_string(),
124            });
125        }
126        self.databases.remove(database_name);
127        Ok(())
128    }
129
130    pub fn list_databases(&self) -> Vec<&Database> {
131        let mut dbs: Vec<&Database> = self.databases.values().collect();
132        dbs.sort_by(|a, b| a.database_name.cmp(&b.database_name));
133        dbs
134    }
135
136    pub fn update_database(
137        &mut self,
138        database_name: &str,
139        kms_key_id: &str,
140    ) -> Result<&Database, TimestreamWriteError> {
141        let db = self.databases.get_mut(database_name).ok_or_else(|| {
142            TimestreamWriteError::DatabaseNotFound {
143                database_name: database_name.to_string(),
144            }
145        })?;
146        db.kms_key_id = Some(kms_key_id.to_string());
147        db.last_updated_time = Utc::now();
148        Ok(db)
149    }
150
151    // --- Table operations ---
152
153    pub fn create_table(
154        &mut self,
155        database_name: &str,
156        table_name: &str,
157        account_id: &str,
158        region: &str,
159        retention_properties: Option<RetentionProperties>,
160        magnetic_store_write_properties: Option<MagneticStoreWriteProperties>,
161    ) -> Result<&Table, TimestreamWriteError> {
162        if !self.databases.contains_key(database_name) {
163            return Err(TimestreamWriteError::DatabaseNotFound {
164                database_name: database_name.to_string(),
165            });
166        }
167
168        let key = (database_name.to_string(), table_name.to_string());
169        if self.tables.contains_key(&key) {
170            return Err(TimestreamWriteError::TableAlreadyExists {
171                database_name: database_name.to_string(),
172                table_name: table_name.to_string(),
173            });
174        }
175
176        let arn = format!(
177            "arn:aws:timestream:{region}:{account_id}:database/{database_name}/table/{table_name}"
178        );
179        let now = Utc::now();
180        let table = Table {
181            table_name: table_name.to_string(),
182            database_name: database_name.to_string(),
183            arn,
184            table_status: "ACTIVE".to_string(),
185            retention_properties: retention_properties.unwrap_or_default(),
186            magnetic_store_write_properties: magnetic_store_write_properties.unwrap_or_default(),
187            creation_time: now,
188            last_updated_time: now,
189            tags: HashMap::new(),
190            records: Vec::new(),
191        };
192
193        self.tables.insert(key.clone(), table);
194
195        // Update table count
196        if let Some(db) = self.databases.get_mut(database_name) {
197            db.table_count += 1;
198        }
199
200        Ok(self.tables.get(&key).unwrap())
201    }
202
203    pub fn describe_table(
204        &self,
205        database_name: &str,
206        table_name: &str,
207    ) -> Result<&Table, TimestreamWriteError> {
208        if !self.databases.contains_key(database_name) {
209            return Err(TimestreamWriteError::DatabaseNotFound {
210                database_name: database_name.to_string(),
211            });
212        }
213        let key = (database_name.to_string(), table_name.to_string());
214        self.tables
215            .get(&key)
216            .ok_or_else(|| TimestreamWriteError::TableNotFound {
217                database_name: database_name.to_string(),
218                table_name: table_name.to_string(),
219            })
220    }
221
222    pub fn delete_table(
223        &mut self,
224        database_name: &str,
225        table_name: &str,
226    ) -> Result<(), TimestreamWriteError> {
227        if !self.databases.contains_key(database_name) {
228            return Err(TimestreamWriteError::DatabaseNotFound {
229                database_name: database_name.to_string(),
230            });
231        }
232        let key = (database_name.to_string(), table_name.to_string());
233        match self.tables.remove(&key) {
234            Some(_) => {
235                if let Some(db) = self.databases.get_mut(database_name) {
236                    db.table_count -= 1;
237                }
238                Ok(())
239            }
240            None => Err(TimestreamWriteError::TableNotFound {
241                database_name: database_name.to_string(),
242                table_name: table_name.to_string(),
243            }),
244        }
245    }
246
247    pub fn list_tables(&self, database_name: &str) -> Result<Vec<&Table>, TimestreamWriteError> {
248        if !self.databases.contains_key(database_name) {
249            return Err(TimestreamWriteError::DatabaseNotFound {
250                database_name: database_name.to_string(),
251            });
252        }
253        let mut tables: Vec<&Table> = self
254            .tables
255            .values()
256            .filter(|t| t.database_name == database_name)
257            .collect();
258        tables.sort_by(|a, b| a.table_name.cmp(&b.table_name));
259        Ok(tables)
260    }
261
262    pub fn update_table(
263        &mut self,
264        database_name: &str,
265        table_name: &str,
266        retention_properties: Option<RetentionProperties>,
267        magnetic_store_write_properties: Option<MagneticStoreWriteProperties>,
268    ) -> Result<&Table, TimestreamWriteError> {
269        if !self.databases.contains_key(database_name) {
270            return Err(TimestreamWriteError::DatabaseNotFound {
271                database_name: database_name.to_string(),
272            });
273        }
274        let key = (database_name.to_string(), table_name.to_string());
275        let table =
276            self.tables
277                .get_mut(&key)
278                .ok_or_else(|| TimestreamWriteError::TableNotFound {
279                    database_name: database_name.to_string(),
280                    table_name: table_name.to_string(),
281                })?;
282
283        if let Some(rp) = retention_properties {
284            table.retention_properties = rp;
285        }
286        if let Some(mswp) = magnetic_store_write_properties {
287            table.magnetic_store_write_properties = mswp;
288        }
289        table.last_updated_time = Utc::now();
290        Ok(table)
291    }
292
293    // --- WriteRecords ---
294
295    pub fn write_records(
296        &mut self,
297        database_name: &str,
298        table_name: &str,
299        records: Vec<Record>,
300    ) -> Result<i32, TimestreamWriteError> {
301        if !self.databases.contains_key(database_name) {
302            return Err(TimestreamWriteError::DatabaseNotFound {
303                database_name: database_name.to_string(),
304            });
305        }
306        let key = (database_name.to_string(), table_name.to_string());
307        let table =
308            self.tables
309                .get_mut(&key)
310                .ok_or_else(|| TimestreamWriteError::TableNotFound {
311                    database_name: database_name.to_string(),
312                    table_name: table_name.to_string(),
313                })?;
314
315        let count = records.len() as i32;
316        table.records.extend(records);
317        Ok(count)
318    }
319
320    // --- Tag operations ---
321
322    pub fn tag_resource(
323        &mut self,
324        resource_arn: &str,
325        tags: HashMap<String, String>,
326    ) -> Result<(), TimestreamWriteError> {
327        // Try databases first
328        for db in self.databases.values_mut() {
329            if db.arn == resource_arn {
330                db.tags.extend(tags);
331                return Ok(());
332            }
333        }
334        // Try tables
335        for table in self.tables.values_mut() {
336            if table.arn == resource_arn {
337                table.tags.extend(tags);
338                return Ok(());
339            }
340        }
341        Err(TimestreamWriteError::ResourceNotFound {
342            resource_arn: resource_arn.to_string(),
343        })
344    }
345
346    pub fn untag_resource(
347        &mut self,
348        resource_arn: &str,
349        tag_keys: &[String],
350    ) -> Result<(), TimestreamWriteError> {
351        // Try databases first
352        for db in self.databases.values_mut() {
353            if db.arn == resource_arn {
354                for key in tag_keys {
355                    db.tags.remove(key);
356                }
357                return Ok(());
358            }
359        }
360        // Try tables
361        for table in self.tables.values_mut() {
362            if table.arn == resource_arn {
363                for key in tag_keys {
364                    table.tags.remove(key);
365                }
366                return Ok(());
367            }
368        }
369        Err(TimestreamWriteError::ResourceNotFound {
370            resource_arn: resource_arn.to_string(),
371        })
372    }
373
374    pub fn list_tags_for_resource(
375        &self,
376        resource_arn: &str,
377    ) -> Result<&HashMap<String, String>, TimestreamWriteError> {
378        for db in self.databases.values() {
379            if db.arn == resource_arn {
380                return Ok(&db.tags);
381            }
382        }
383        for table in self.tables.values() {
384            if table.arn == resource_arn {
385                return Ok(&table.tags);
386            }
387        }
388        Err(TimestreamWriteError::ResourceNotFound {
389            resource_arn: resource_arn.to_string(),
390        })
391    }
392
393    // --- BatchLoadTask operations ---
394
395    pub fn create_batch_load_task(
396        &mut self,
397        task_id: String,
398        database_name: &str,
399        table_name: &str,
400        data_source_configuration: Value,
401        report_configuration: Value,
402        data_model_configuration: Option<Value>,
403    ) -> &BatchLoadTask {
404        let now = Utc::now();
405        let task = BatchLoadTask {
406            task_id: task_id.clone(),
407            task_status: "CREATED".to_string(),
408            database_name: database_name.to_string(),
409            table_name: table_name.to_string(),
410            data_source_configuration,
411            report_configuration,
412            data_model_configuration,
413            creation_time: now,
414            last_updated_time: now,
415            error_message: None,
416        };
417        self.batch_load_tasks.insert(task_id.clone(), task);
418        self.batch_load_tasks.get(&task_id).unwrap()
419    }
420
421    pub fn describe_batch_load_task(
422        &self,
423        task_id: &str,
424    ) -> Result<&BatchLoadTask, TimestreamWriteError> {
425        self.batch_load_tasks
426            .get(task_id)
427            .ok_or_else(|| TimestreamWriteError::TaskNotFound {
428                task_id: task_id.to_string(),
429            })
430    }
431
432    pub fn list_batch_load_tasks(&self, task_status_filter: Option<&str>) -> Vec<&BatchLoadTask> {
433        let mut tasks: Vec<&BatchLoadTask> = self
434            .batch_load_tasks
435            .values()
436            .filter(|t| {
437                task_status_filter
438                    .map(|s| t.task_status == s)
439                    .unwrap_or(true)
440            })
441            .collect();
442        tasks.sort_by(|a, b| a.task_id.cmp(&b.task_id));
443        tasks
444    }
445
446    pub fn resume_batch_load_task(&mut self, task_id: &str) -> Result<(), TimestreamWriteError> {
447        let task = self.batch_load_tasks.get_mut(task_id).ok_or_else(|| {
448            TimestreamWriteError::TaskNotFound {
449                task_id: task_id.to_string(),
450            }
451        })?;
452        if task.task_status != "FAILED" {
453            return Err(TimestreamWriteError::TaskCannotBeResumed {
454                task_id: task_id.to_string(),
455                task_status: task.task_status.clone(),
456            });
457        }
458        task.task_status = "IN_PROGRESS".to_string();
459        task.last_updated_time = Utc::now();
460        Ok(())
461    }
462}
winterbaume_timestreamwrite/state.rs

winterbaume_timestreamwrite/
state.rs
