1use chrono::{DateTime, Utc};
2use parking_lot::RwLock;
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::collections::{BTreeMap, HashMap};
6use std::sync::Arc;
7
8fn empty_stream_records() -> Arc<RwLock<Vec<StreamRecord>>> {
9 Arc::new(RwLock::new(Vec::new()))
10}
11
12mod stream_records_serde {
17 use super::{Arc, RwLock, StreamRecord};
18 use serde::{Deserialize, Deserializer, Serialize, Serializer};
19
20 pub fn serialize<S: Serializer>(
21 v: &Arc<RwLock<Vec<StreamRecord>>>,
22 s: S,
23 ) -> Result<S::Ok, S::Error> {
24 v.read().serialize(s)
25 }
26
27 pub fn deserialize<'de, D: Deserializer<'de>>(
28 d: D,
29 ) -> Result<Arc<RwLock<Vec<StreamRecord>>>, D::Error> {
30 let records = Vec::<StreamRecord>::deserialize(d)?;
31 for r in &records {
35 crate::streams::observe_stream_sequence(&r.dynamodb.sequence_number);
36 }
37 Ok(Arc::new(RwLock::new(records)))
38 }
39}
40
41pub type AttributeValue = Value;
44
45pub fn attribute_type_and_value(av: &Value) -> Option<(&str, &Value)> {
48 let obj = av.as_object()?;
49 if obj.len() != 1 {
50 return None;
51 }
52 let (k, v) = obj.iter().next()?;
53 Some((k.as_str(), v))
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct KeySchemaElement {
58 pub attribute_name: String,
59 pub key_type: String, }
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct AttributeDefinition {
64 pub attribute_name: String,
65 pub attribute_type: String, }
67
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct ProvisionedThroughput {
70 pub read_capacity_units: i64,
71 pub write_capacity_units: i64,
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct OnDemandThroughput {
80 pub max_read_request_units: i64,
81 pub max_write_request_units: i64,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct GlobalSecondaryIndex {
86 pub index_name: String,
87 pub key_schema: Vec<KeySchemaElement>,
88 pub projection: Projection,
89 pub provisioned_throughput: Option<ProvisionedThroughput>,
90 pub on_demand_throughput: Option<OnDemandThroughput>,
91}
92
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct LocalSecondaryIndex {
95 pub index_name: String,
96 pub key_schema: Vec<KeySchemaElement>,
97 pub projection: Projection,
98}
99
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct Projection {
102 pub projection_type: String, pub non_key_attributes: Vec<String>,
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct DynamoTable {
108 pub name: String,
109 pub arn: String,
110 pub table_id: String,
111 pub key_schema: Vec<KeySchemaElement>,
112 pub attribute_definitions: Vec<AttributeDefinition>,
113 pub provisioned_throughput: ProvisionedThroughput,
114 pub items: Vec<HashMap<String, AttributeValue>>,
115 pub gsi: Vec<GlobalSecondaryIndex>,
116 pub lsi: Vec<LocalSecondaryIndex>,
117 pub tags: BTreeMap<String, String>,
118 pub created_at: DateTime<Utc>,
119 pub status: String,
120 pub item_count: i64,
121 pub size_bytes: i64,
122 pub billing_mode: String, pub ttl_attribute: Option<String>,
124 pub ttl_enabled: bool,
125 pub resource_policy: Option<String>,
126 pub pitr_enabled: bool,
128 pub kinesis_destinations: Vec<KinesisDestination>,
130 pub contributor_insights_status: String,
132 pub contributor_insights_counters: BTreeMap<String, u64>,
134 pub stream_enabled: bool,
136 pub stream_view_type: Option<String>, pub stream_arn: Option<String>,
138 #[serde(with = "stream_records_serde", default = "empty_stream_records")]
141 pub stream_records: Arc<RwLock<Vec<StreamRecord>>>,
142 pub sse_type: Option<String>,
144 pub sse_kms_key_arn: Option<String>,
146 pub deletion_protection_enabled: bool,
150 pub on_demand_throughput: Option<OnDemandThroughput>,
154}
155
156#[derive(Debug, Clone, Serialize, Deserialize)]
157pub struct StreamRecord {
158 pub event_id: String,
159 pub event_name: String, pub event_version: String,
161 pub event_source: String,
162 pub aws_region: String,
163 pub dynamodb: DynamoDbStreamRecord,
164 pub event_source_arn: String,
165 pub timestamp: DateTime<Utc>,
166}
167
168#[derive(Debug, Clone, Serialize, Deserialize)]
169pub struct DynamoDbStreamRecord {
170 pub keys: HashMap<String, AttributeValue>,
171 pub new_image: Option<HashMap<String, AttributeValue>>,
172 pub old_image: Option<HashMap<String, AttributeValue>>,
173 pub sequence_number: String,
174 pub size_bytes: i64,
175 pub stream_view_type: String,
176}
177
178#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct KinesisDestination {
180 pub stream_arn: String,
181 pub destination_status: String,
182 pub approximate_creation_date_time_precision: String,
183}
184
185#[derive(Debug, Clone, Serialize, Deserialize)]
186pub struct BackupDescription {
187 pub backup_arn: String,
188 pub backup_name: String,
189 pub table_name: String,
190 pub table_arn: String,
191 pub backup_status: String,
192 pub backup_type: String,
193 pub backup_creation_date: DateTime<Utc>,
194 pub key_schema: Vec<KeySchemaElement>,
195 pub attribute_definitions: Vec<AttributeDefinition>,
196 pub provisioned_throughput: ProvisionedThroughput,
197 pub billing_mode: String,
198 pub item_count: i64,
199 pub size_bytes: i64,
200 pub items: Vec<HashMap<String, AttributeValue>>,
202 #[serde(default)]
207 pub gsi: Vec<GlobalSecondaryIndex>,
208 #[serde(default)]
209 pub lsi: Vec<LocalSecondaryIndex>,
210 #[serde(default)]
211 pub tags: BTreeMap<String, String>,
212 #[serde(default)]
213 pub ttl_attribute: Option<String>,
214 #[serde(default)]
215 pub ttl_enabled: bool,
216 #[serde(default)]
217 pub sse_type: Option<String>,
218 #[serde(default)]
219 pub sse_kms_key_arn: Option<String>,
220 #[serde(default)]
221 pub stream_enabled: bool,
222 #[serde(default)]
223 pub stream_view_type: Option<String>,
224}
225
226#[derive(Debug, Clone, Serialize, Deserialize)]
227pub struct GlobalTableDescription {
228 pub global_table_name: String,
229 pub global_table_arn: String,
230 pub global_table_status: String,
231 pub creation_date: DateTime<Utc>,
232 pub replication_group: Vec<ReplicaDescription>,
233}
234
235#[derive(Debug, Clone, Serialize, Deserialize)]
236pub struct ReplicaDescription {
237 pub region_name: String,
238 pub replica_status: String,
239}
240
241#[derive(Debug, Clone, Serialize, Deserialize)]
242pub struct ExportDescription {
243 pub export_arn: String,
244 pub export_status: String,
245 pub table_arn: String,
246 pub s3_bucket: String,
247 pub s3_prefix: Option<String>,
248 pub export_format: String,
249 pub start_time: DateTime<Utc>,
250 pub end_time: DateTime<Utc>,
251 pub export_time: DateTime<Utc>,
252 pub item_count: i64,
253 pub billed_size_bytes: i64,
254}
255
256#[derive(Debug, Clone, Serialize, Deserialize)]
257pub struct ImportDescription {
258 pub import_arn: String,
259 pub import_status: String,
260 pub table_arn: String,
261 pub table_name: String,
262 pub s3_bucket_source: String,
263 pub input_format: String,
264 pub start_time: DateTime<Utc>,
265 pub end_time: DateTime<Utc>,
266 pub processed_item_count: i64,
267 pub processed_size_bytes: i64,
268}
269
270impl DynamoTable {
271 pub fn hash_key_name(&self) -> &str {
273 self.key_schema
274 .iter()
275 .find(|k| k.key_type == "HASH")
276 .map(|k| k.attribute_name.as_str())
277 .unwrap_or("")
278 }
279
280 pub fn range_key_name(&self) -> Option<&str> {
282 self.key_schema
283 .iter()
284 .find(|k| k.key_type == "RANGE")
285 .map(|k| k.attribute_name.as_str())
286 }
287
288 pub fn find_item_index(&self, key: &HashMap<String, AttributeValue>) -> Option<usize> {
290 let hash_key = self.hash_key_name();
291 let range_key = self.range_key_name();
292
293 self.items.iter().position(|item| {
294 let hash_match = match (item.get(hash_key), key.get(hash_key)) {
295 (Some(a), Some(b)) => a == b,
296 _ => false,
297 };
298 if !hash_match {
299 return false;
300 }
301 match range_key {
302 Some(rk) => match (item.get(rk), key.get(rk)) {
303 (Some(a), Some(b)) => a == b,
304 (None, None) => true,
305 _ => false,
306 },
307 None => true,
308 }
309 })
310 }
311
312 fn estimate_item_size(item: &HashMap<String, AttributeValue>) -> i64 {
314 let mut size: i64 = 0;
315 for (k, v) in item {
316 size += k.len() as i64;
317 size += Self::estimate_value_size(v);
318 }
319 size
320 }
321
322 fn estimate_value_size(v: &Value) -> i64 {
323 match v {
324 Value::Object(obj) => {
325 if let Some(s) = obj.get("S").and_then(|v| v.as_str()) {
326 s.len() as i64
327 } else if let Some(n) = obj.get("N").and_then(|v| v.as_str()) {
328 n.len() as i64
329 } else if obj.contains_key("BOOL") || obj.contains_key("NULL") {
330 1
331 } else if let Some(l) = obj.get("L").and_then(|v| v.as_array()) {
332 3 + l.iter().map(Self::estimate_value_size).sum::<i64>()
333 } else if let Some(m) = obj.get("M").and_then(|v| v.as_object()) {
334 3 + m
335 .iter()
336 .map(|(k, v)| k.len() as i64 + Self::estimate_value_size(v))
337 .sum::<i64>()
338 } else if let Some(ss) = obj.get("SS").and_then(|v| v.as_array()) {
339 ss.iter()
340 .filter_map(|v| v.as_str())
341 .map(|s| s.len() as i64)
342 .sum()
343 } else if let Some(ns) = obj.get("NS").and_then(|v| v.as_array()) {
344 ns.iter()
345 .filter_map(|v| v.as_str())
346 .map(|s| s.len() as i64)
347 .sum()
348 } else if let Some(b) = obj.get("B").and_then(|v| v.as_str()) {
349 (b.len() as i64 * 3) / 4
351 } else {
352 v.to_string().len() as i64
353 }
354 }
355 _ => v.to_string().len() as i64,
356 }
357 }
358
359 pub fn record_key_access(&mut self, key: &HashMap<String, AttributeValue>) {
362 if self.contributor_insights_status != "ENABLED" {
363 return;
364 }
365 let hash_key = self.hash_key_name().to_string();
366 if let Some(pk_value) = key.get(&hash_key) {
367 let key_str = pk_value.to_string();
368 *self
369 .contributor_insights_counters
370 .entry(key_str)
371 .or_insert(0) += 1;
372 }
373 }
374
375 pub fn record_item_access(&mut self, item: &HashMap<String, AttributeValue>) {
377 if self.contributor_insights_status != "ENABLED" {
378 return;
379 }
380 let hash_key = self.hash_key_name().to_string();
381 if let Some(pk_value) = item.get(&hash_key) {
382 let key_str = pk_value.to_string();
383 *self
384 .contributor_insights_counters
385 .entry(key_str)
386 .or_insert(0) += 1;
387 }
388 }
389
390 pub fn top_contributors(&self, n: usize) -> Vec<(&str, u64)> {
392 let mut entries: Vec<(&str, u64)> = self
393 .contributor_insights_counters
394 .iter()
395 .map(|(k, &v)| (k.as_str(), v))
396 .collect();
397 entries.sort_by_key(|e| std::cmp::Reverse(e.1));
398 entries.truncate(n);
399 entries
400 }
401
402 pub fn recalculate_stats(&mut self) {
404 self.item_count = self.items.len() as i64;
405 self.size_bytes = self.items.iter().map(Self::estimate_item_size).sum::<i64>();
406 }
407}
408
409#[derive(Debug, Clone, Serialize, Deserialize)]
410pub struct DynamoDbState {
411 pub account_id: String,
412 pub region: String,
413 pub tables: BTreeMap<String, DynamoTable>,
414 pub backups: BTreeMap<String, BackupDescription>,
415 pub global_tables: BTreeMap<String, GlobalTableDescription>,
416 pub exports: BTreeMap<String, ExportDescription>,
417 pub imports: BTreeMap<String, ImportDescription>,
418}
419
420#[derive(Debug, Clone, Serialize, Deserialize)]
424pub struct DynamoDbSnapshot {
425 pub schema_version: u32,
426 #[serde(default)]
428 pub accounts: Option<fakecloud_core::multi_account::MultiAccountState<DynamoDbState>>,
429 #[serde(default)]
431 pub state: Option<DynamoDbState>,
432}
433
434pub const DYNAMODB_SNAPSHOT_SCHEMA_VERSION: u32 = 2;
435
436impl DynamoDbState {
437 pub fn new(account_id: &str, region: &str) -> Self {
438 Self {
439 account_id: account_id.to_string(),
440 region: region.to_string(),
441 tables: BTreeMap::new(),
442 backups: BTreeMap::new(),
443 global_tables: BTreeMap::new(),
444 exports: BTreeMap::new(),
445 imports: BTreeMap::new(),
446 }
447 }
448
449 pub fn reset(&mut self) {
450 self.tables.clear();
451 self.backups.clear();
452 self.global_tables.clear();
453 self.exports.clear();
454 self.imports.clear();
455 }
456}
457
458impl fakecloud_core::multi_account::AccountState for DynamoDbState {
459 fn new_for_account(account_id: &str, region: &str, _endpoint: &str) -> Self {
460 Self::new(account_id, region)
461 }
462}
463
464pub type SharedDynamoDbState =
465 Arc<RwLock<fakecloud_core::multi_account::MultiAccountState<DynamoDbState>>>;
466
467#[cfg(test)]
468mod tests {
469 use super::*;
470 use serde_json::json;
471
472 #[test]
473 fn attribute_type_and_value_valid() {
474 let v = json!({"S": "hi"});
475 let (ty, val) = attribute_type_and_value(&v).unwrap();
476 assert_eq!(ty, "S");
477 assert_eq!(val, &json!("hi"));
478 }
479
480 #[test]
481 fn attribute_type_and_value_empty_returns_none() {
482 let v = json!({});
483 assert!(attribute_type_and_value(&v).is_none());
484 }
485
486 #[test]
487 fn attribute_type_and_value_multiple_entries_returns_none() {
488 let v = json!({"S": "hi", "N": "1"});
489 assert!(attribute_type_and_value(&v).is_none());
490 }
491
492 #[test]
493 fn attribute_type_and_value_non_object_returns_none() {
494 let v = json!("not-object");
495 assert!(attribute_type_and_value(&v).is_none());
496 }
497
498 #[test]
499 fn account_state_trait_impl() {
500 use fakecloud_core::multi_account::AccountState;
501 let state = DynamoDbState::new_for_account("123", "us-east-1", "");
502 assert_eq!(state.account_id, "123");
503 assert_eq!(state.region, "us-east-1");
504 }
505
506 #[test]
507 fn new_and_reset() {
508 let state = DynamoDbState::new("123", "us-east-1");
509 assert!(state.tables.is_empty());
510 }
511
512 fn table_with_hash_key(hash: &str) -> DynamoTable {
513 DynamoTable {
514 name: "t".to_string(),
515 arn: "arn:aws:dynamodb:us-east-1:123:table/t".to_string(),
516 table_id: "id".to_string(),
517 key_schema: vec![KeySchemaElement {
518 attribute_name: hash.to_string(),
519 key_type: "HASH".to_string(),
520 }],
521 attribute_definitions: vec![],
522 provisioned_throughput: ProvisionedThroughput {
523 read_capacity_units: 1,
524 write_capacity_units: 1,
525 },
526 items: Vec::new(),
527 gsi: Vec::new(),
528 lsi: Vec::new(),
529 tags: BTreeMap::new(),
530 created_at: Utc::now(),
531 status: "ACTIVE".to_string(),
532 item_count: 0,
533 size_bytes: 0,
534 billing_mode: "PROVISIONED".to_string(),
535 ttl_attribute: None,
536 ttl_enabled: false,
537 resource_policy: None,
538 pitr_enabled: false,
539 kinesis_destinations: Vec::new(),
540 contributor_insights_status: "DISABLED".to_string(),
541 contributor_insights_counters: BTreeMap::new(),
542 stream_enabled: false,
543 stream_view_type: None,
544 stream_arn: None,
545 stream_records: empty_stream_records(),
546 sse_type: None,
547 sse_kms_key_arn: None,
548 deletion_protection_enabled: false,
549 on_demand_throughput: None,
550 }
551 }
552
553 #[test]
554 fn hash_key_name_extracts_from_schema() {
555 let t = table_with_hash_key("pk");
556 assert_eq!(t.hash_key_name(), "pk");
557 }
558
559 #[test]
560 fn hash_key_name_empty_when_no_hash_schema() {
561 let mut t = table_with_hash_key("pk");
562 t.key_schema.clear();
563 assert_eq!(t.hash_key_name(), "");
564 }
565
566 #[test]
567 fn record_key_access_noop_when_disabled() {
568 let mut t = table_with_hash_key("pk");
569 let mut key = HashMap::new();
570 key.insert("pk".to_string(), json!({"S": "a"}));
571 t.record_key_access(&key);
572 assert!(t.contributor_insights_counters.is_empty());
573 }
574
575 #[test]
576 fn record_key_access_increments_when_enabled() {
577 let mut t = table_with_hash_key("pk");
578 t.contributor_insights_status = "ENABLED".to_string();
579 let mut key = HashMap::new();
580 key.insert("pk".to_string(), json!({"S": "a"}));
581 t.record_key_access(&key);
582 t.record_key_access(&key);
583 assert_eq!(t.contributor_insights_counters.values().sum::<u64>(), 2);
584 }
585
586 #[test]
587 fn record_item_access_uses_hash_key_from_item() {
588 let mut t = table_with_hash_key("pk");
589 t.contributor_insights_status = "ENABLED".to_string();
590 let mut item = HashMap::new();
591 item.insert("pk".to_string(), json!({"S": "user-1"}));
592 item.insert("other".to_string(), json!({"N": "42"}));
593 t.record_item_access(&item);
594 assert_eq!(t.contributor_insights_counters.values().sum::<u64>(), 1);
595 }
596
597 #[test]
598 fn top_contributors_returns_sorted() {
599 let mut t = table_with_hash_key("pk");
600 t.contributor_insights_counters.insert("a".to_string(), 3);
601 t.contributor_insights_counters.insert("b".to_string(), 10);
602 t.contributor_insights_counters.insert("c".to_string(), 1);
603 let top = t.top_contributors(2);
604 assert_eq!(top.len(), 2);
605 assert_eq!(top[0], ("b", 10));
606 assert_eq!(top[1], ("a", 3));
607 }
608
609 #[test]
610 fn recalculate_stats_matches_items() {
611 let mut t = table_with_hash_key("pk");
612 let mut item1 = HashMap::new();
613 item1.insert("pk".to_string(), json!({"S": "hello"}));
614 let mut item2 = HashMap::new();
615 item2.insert("pk".to_string(), json!({"N": "42"}));
616 item2.insert("flag".to_string(), json!({"BOOL": true}));
617 t.items.push(item1);
618 t.items.push(item2);
619 t.recalculate_stats();
620 assert_eq!(t.item_count, 2);
621 assert!(t.size_bytes > 0);
622 }
623
624 #[test]
625 fn estimate_value_size_covers_all_types() {
626 let s = DynamoTable::estimate_value_size(&json!({"S": "abc"}));
627 assert_eq!(s, 3);
628 let n = DynamoTable::estimate_value_size(&json!({"N": "42"}));
629 assert_eq!(n, 2);
630 let b = DynamoTable::estimate_value_size(&json!({"BOOL": true}));
631 assert_eq!(b, 1);
632 let null = DynamoTable::estimate_value_size(&json!({"NULL": true}));
633 assert_eq!(null, 1);
634 let l = DynamoTable::estimate_value_size(&json!({"L": [{"S": "x"}, {"S": "yy"}]}));
635 assert_eq!(l, 6);
636 let m = DynamoTable::estimate_value_size(&json!({"M": {"key": {"S": "v"}}}));
637 assert_eq!(m, 7);
638 let ss = DynamoTable::estimate_value_size(&json!({"SS": ["ab", "cde"]}));
639 assert_eq!(ss, 5);
640 let ns = DynamoTable::estimate_value_size(&json!({"NS": ["12", "345"]}));
641 assert_eq!(ns, 5);
642 let bin = DynamoTable::estimate_value_size(&json!({"B": "AAAAAAAA"}));
643 assert_eq!(bin, 6);
644 }
645}