Skip to main content

fakecloud_dynamodb/
export_import.rs

1//! Startup bulk-importer for AWS-format DynamoDB S3 exports (DYNAMODB_JSON).
2//!
3//! Reads a local copy of an AWS DynamoDB S3 export (manifests + gzipped data
4//! files) plus the table's `describe-table` JSON, and materialises it as a
5//! fresh table in state. The AWS attribute wire format IS our storage format
6//! (`AttributeValue` is a `serde_json::Value`), so items copy through verbatim
7//! with no type decoding.
8
9use std::collections::HashMap;
10use std::io::Read as _;
11use std::path::{Path, PathBuf};
12
13use serde_json::Value;
14
15use crate::state::{DynamoTable, ProvisionedThroughput, SharedDynamoDbState};
16
17/// A single DynamoDB item: attribute name -> typed AWS wire value.
18type Item = HashMap<String, Value>;
19
20/// Outcome of a startup export import.
21///
22/// The import is idempotent: in persistent-storage mode the snapshot is loaded
23/// before this import runs, so on a restart with the same import flags the
24/// target table is already present. Rather than erroring (which would refuse to
25/// boot), the import is skipped and the already-loaded table is left untouched.
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub enum ImportOutcome {
28    /// A fresh table was materialised from the export with `items` items.
29    Imported { table: String, items: usize },
30    /// A table of the same name already existed in state; import was skipped and
31    /// the existing data was left untouched.
32    SkippedExisting { table: String },
33}
34
35/// Import an AWS-format DynamoDB export as a new table.
36///
37/// Idempotent: if the target table already exists in state (for example loaded
38/// from a persisted snapshot on restart), the import is skipped and the existing
39/// data is preserved instead of erroring out.
40pub fn import_aws_export(
41    state: &SharedDynamoDbState,
42    account_id: &str,
43    region: &str,
44    export_dir: &Path,      // folder holding manifest-summary.json etc.
45    describe_table: &Value, // parsed `aws dynamodb describe-table` JSON
46) -> Result<ImportOutcome, String> {
47    // Accept either the full `{"Table": {...}}` dump or a bare table object.
48    let shape = describe_table.get("Table").unwrap_or(describe_table);
49    let table_name = shape["TableName"]
50        .as_str()
51        .ok_or("describe-table: TableName missing or not a string")?
52        .to_string();
53
54    // Idempotent restart guard: if the table is already present (e.g. restored
55    // from a persisted snapshot before this import runs), leave it untouched and
56    // skip rather than error and refuse to boot. Checked before reading the
57    // export so a no-op restart does no work.
58    if table_exists(state, account_id, &table_name) {
59        tracing::warn!(
60            table = %table_name,
61            "skipping DynamoDB export import: table already exists in state; existing data left untouched"
62        );
63        return Ok(ImportOutcome::SkippedExisting { table: table_name });
64    }
65
66    let (items, declared_total) = read_export_items(export_dir)?;
67    let item_count = items.len();
68    // Manifest integrity: the summary's declared grand total must match what was
69    // actually read. Absent field -> skip the check (older/partial manifests).
70    if let Some(declared) = declared_total {
71        if declared != item_count {
72            return Err(format!(
73                "table {table_name}: manifest-summary declares itemCount {declared} but {item_count} items were read (truncated or corrupt export)"
74            ));
75        }
76    }
77    let table = build_table(&table_name, region, account_id, shape, items)?;
78
79    let mut guard = state.write();
80    let account = guard.get_or_create(account_id);
81    // Re-check under the write lock to close any check-then-insert gap and to
82    // handle a genuinely conflicting table gracefully (warn + skip, never panic).
83    if account.tables.contains_key(&table_name) {
84        tracing::warn!(
85            table = %table_name,
86            "skipping DynamoDB export import: table already exists in state; existing data left untouched"
87        );
88        return Ok(ImportOutcome::SkippedExisting { table: table_name });
89    }
90    account.tables.insert(table_name.clone(), table);
91
92    Ok(ImportOutcome::Imported {
93        table: table_name,
94        items: item_count,
95    })
96}
97
98/// Multi-table counterpart to `import_aws_export`: imports every immediate
99/// subdirectory of `root_dir` as its own table (see docs/services/dynamodb.md
100/// for the on-disk layout). Stops at the first table that fails; tables
101/// already imported before that point are not rolled back.
102pub fn import_aws_exports_dir(
103    state: &SharedDynamoDbState,
104    account_id: &str,
105    region: &str,
106    root_dir: &Path,
107) -> Result<Vec<ImportOutcome>, String> {
108    let mut subdirs: Vec<PathBuf> = std::fs::read_dir(root_dir)
109        .map_err(|e| format!("read {}: {e}", root_dir.display()))?
110        .filter_map(|entry| entry.ok())
111        .map(|entry| entry.path())
112        .filter(|path| path.is_dir())
113        .collect();
114    // Deterministic order: stable logging/import order across runs regardless
115    // of the OS's directory-listing order.
116    subdirs.sort();
117
118    if subdirs.is_empty() {
119        return Err(format!(
120            "no per-table subdirectories found under {}",
121            root_dir.display()
122        ));
123    }
124
125    let mut outcomes = Vec::with_capacity(subdirs.len());
126    for dir in subdirs {
127        let describe = read_json(&dir.join("describe-table.json"))
128            .map_err(|e| format!("{}: {e}", dir.display()))?;
129        let outcome = import_aws_export(state, account_id, region, &dir, &describe)
130            .map_err(|e| format!("{}: {e}", dir.display()))?;
131        outcomes.push(outcome);
132    }
133
134    Ok(outcomes)
135}
136
137/// True if `account_id` already holds a table named `table_name`.
138fn table_exists(state: &SharedDynamoDbState, account_id: &str, table_name: &str) -> bool {
139    state
140        .read()
141        .get(account_id)
142        .is_some_and(|account| account.tables.contains_key(table_name))
143}
144
145/// Walk the manifests (not a prefix listing) and read every item. Returns the
146/// items plus the summary manifest's declared `itemCount` (if present) so the
147/// caller can verify the grand total.
148fn read_export_items(export_dir: &Path) -> Result<(Vec<Item>, Option<usize>), String> {
149    let summary = read_json(&export_dir.join("manifest-summary.json"))?;
150    let format = summary["exportFormat"].as_str().unwrap_or_default();
151    if format != "DYNAMODB_JSON" {
152        return Err(format!(
153            "unsupported exportFormat {format:?} (only DYNAMODB_JSON is supported)"
154        ));
155    }
156
157    // `manifest-files.json` is JSON Lines (one object per line), each naming a
158    // data file via `dataFileS3Key`.
159    let manifest = read_text(&export_dir.join("manifest-files.json"))?;
160    let mut items = Vec::new();
161    for line in nonempty_lines(&manifest) {
162        let entry: Value =
163            serde_json::from_str(line).map_err(|e| format!("parse manifest-files line: {e}"))?;
164        let s3_key = entry["dataFileS3Key"]
165            .as_str()
166            .ok_or("manifest-files: dataFileS3Key missing")?;
167        // The manifest key uses the standard `AWSDynamoDB/{export-id}/data/<file>`
168        // layout; resolve it against the local export as `<export_dir>/data/<basename>`.
169        let basename = s3_key.rsplit('/').next().unwrap_or(s3_key);
170        let file_items = read_data_file(&export_dir.join("data").join(basename))?;
171        // Per-file integrity: when the manifest entry declares an itemCount it
172        // must match the records actually read from that data file.
173        if let Some(declared) = entry["itemCount"].as_u64() {
174            if declared as usize != file_items.len() {
175                return Err(format!(
176                    "data file {basename}: manifest declares itemCount {declared} but {} records were read (truncated or corrupt export)",
177                    file_items.len()
178                ));
179            }
180        }
181        items.extend(file_items);
182    }
183    let declared_total = summary["itemCount"].as_u64().map(|n| n as usize);
184    Ok((items, declared_total))
185}
186
187/// Read one gzipped data file; each line is a `{"Item": {...}}` record.
188fn read_data_file(path: &Path) -> Result<Vec<Item>, String> {
189    let file = std::fs::File::open(path).map_err(|e| format!("open {}: {e}", path.display()))?;
190    let mut contents = String::new();
191    flate2::read::GzDecoder::new(file)
192        .read_to_string(&mut contents)
193        .map_err(|e| format!("gunzip {}: {e}", path.display()))?;
194
195    // No dedup: AWS exports carry unique primary keys, so a single data file
196    // never contains a duplicate. Colliding keys would both be retained.
197    nonempty_lines(&contents)
198        .map(|line| {
199            let record: Value =
200                serde_json::from_str(line).map_err(|e| format!("parse data line: {e}"))?;
201            record["Item"]
202                .as_object()
203                .map(|obj| obj.clone().into_iter().collect())
204                .ok_or_else(|| "data line missing Item object".to_string())
205        })
206        .collect()
207}
208
209/// Build a fresh `DynamoTable` from a `describe-table` shape and its items.
210fn build_table(
211    name: &str,
212    region: &str,
213    account_id: &str,
214    shape: &Value,
215    items: Vec<Item>,
216) -> Result<DynamoTable, String> {
217    let key_schema =
218        crate::parse_key_schema(&shape["KeySchema"]).map_err(|e| format!("KeySchema: {e}"))?;
219    let attribute_definitions = crate::parse_attribute_definitions(&shape["AttributeDefinitions"])
220        .map_err(|e| format!("AttributeDefinitions: {e}"))?;
221
222    // Guard against a wrong/hand-edited describe-table producing an
223    // ACTIVE-but-key-corrupt table. For each declared key attribute (HASH and,
224    // if present, RANGE) every item must (1) carry the attribute and (2) carry
225    // it with the scalar type declared in AttributeDefinitions -- the same
226    // presence + type parity the normal write path enforces
227    // (validate_key_in_item / check_key_type). A wrong-typed key would let a
228    // correctly-typed read fail to find the row, so the data would appear to
229    // vanish. Fail the whole import so no partial table lands.
230    for elem in &key_schema {
231        if elem.key_type != "HASH" && elem.key_type != "RANGE" {
232            continue;
233        }
234        let role = if elem.key_type == "HASH" {
235            "partition"
236        } else {
237            "sort"
238        };
239        let expected_type = attribute_definitions
240            .iter()
241            .find(|d| d.attribute_name == elem.attribute_name)
242            .map(|d| d.attribute_type.as_str());
243        for item in &items {
244            let Some(value) = item.get(&elem.attribute_name) else {
245                return Err(format!(
246                    "table {name}: item missing {role} key attribute {:?} declared in KeySchema: {item:?}",
247                    elem.attribute_name
248                ));
249            };
250            if let Some(expected) = expected_type {
251                let actual = crate::state::attribute_type_and_value(value).map(|(ty, _)| ty);
252                if actual != Some(expected) {
253                    return Err(format!(
254                        "table {name}: {role} key attribute {:?} has type {} but AttributeDefinitions declares {expected}: {item:?}",
255                        elem.attribute_name,
256                        actual.unwrap_or("<none>"),
257                    ));
258                }
259            }
260        }
261    }
262
263    // BillingModeSummary wins over a bare BillingMode; default to PROVISIONED.
264    let billing_mode = shape["BillingModeSummary"]["BillingMode"]
265        .as_str()
266        .or_else(|| shape["BillingMode"].as_str())
267        .unwrap_or("PROVISIONED")
268        .to_string();
269    let provisioned_throughput = if billing_mode == "PAY_PER_REQUEST" {
270        ProvisionedThroughput {
271            read_capacity_units: 0,
272            write_capacity_units: 0,
273        }
274    } else {
275        crate::parse_provisioned_throughput(&shape["ProvisionedThroughput"])
276            .map_err(|e| format!("ProvisionedThroughput: {e}"))?
277    };
278
279    let mut table = DynamoTable {
280        name: name.to_string(),
281        arn: format!("arn:aws:dynamodb:{region}:{account_id}:table/{name}"),
282        table_id: uuid::Uuid::new_v4().to_string().replace('-', ""),
283        key_schema,
284        attribute_definitions,
285        provisioned_throughput,
286        items,
287        gsi: crate::parse_gsi(&shape["GlobalSecondaryIndexes"], &billing_mode),
288        lsi: crate::parse_lsi(&shape["LocalSecondaryIndexes"]),
289        tags: crate::parse_tags(&shape["Tags"]),
290        created_at: chrono::Utc::now(),
291        status: "ACTIVE".to_string(),
292        item_count: 0,
293        size_bytes: 0,
294        billing_mode,
295        ttl_attribute: None,
296        ttl_enabled: false,
297        resource_policy: None,
298        pitr_enabled: false,
299        kinesis_destinations: Vec::new(),
300        contributor_insights_status: "DISABLED".to_string(),
301        contributor_insights_counters: std::collections::BTreeMap::new(),
302        stream_enabled: false,
303        stream_view_type: None,
304        stream_arn: None,
305        stream_records: std::sync::Arc::new(parking_lot::RwLock::new(Vec::new())),
306        sse_type: None,
307        sse_kms_key_arn: None,
308        deletion_protection_enabled: false,
309        on_demand_throughput: None,
310        table_class: "STANDARD".to_string(),
311    };
312    table.recalculate_stats(); // fills item_count / size_bytes
313    Ok(table)
314}
315
316fn read_text(path: &Path) -> Result<String, String> {
317    std::fs::read_to_string(path).map_err(|e| format!("read {}: {e}", path.display()))
318}
319
320fn read_json(path: &Path) -> Result<Value, String> {
321    serde_json::from_str(&read_text(path)?).map_err(|e| format!("parse {}: {e}", path.display()))
322}
323
324fn nonempty_lines(text: &str) -> impl Iterator<Item = &str> {
325    text.lines().filter(|l| !l.trim().is_empty())
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331    use serde_json::json;
332    use std::path::PathBuf;
333
334    #[test]
335    fn imports_aws_export_into_state() {
336        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
337            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
338        ));
339
340        let fixtures = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
341        let describe = read_json(&fixtures.join("describe-table.json")).unwrap();
342        let export_dir = fixtures.join("export");
343
344        let outcome =
345            import_aws_export(&state, "123456789012", "us-east-1", &export_dir, &describe)
346                .expect("import should succeed");
347        assert_eq!(
348            outcome,
349            ImportOutcome::Imported {
350                table: "Music".to_string(),
351                items: 3,
352            }
353        );
354
355        // Table lives in state with the right item count.
356        let mut guard = state.write();
357        let table = guard
358            .get_or_create("123456789012")
359            .tables
360            .get("Music")
361            .expect("table in state");
362        assert_eq!(table.items.len(), 3);
363        assert_eq!(table.item_count, 3);
364
365        // Type preservation: the "Foo" item carries every DynamoDB attribute
366        // type and each one survives verbatim. Complex types are asserted for
367        // FULL equality (not spot-checked) so a serializer regression that
368        // dropped a later list element / set member would be caught.
369        let foo = table
370            .items
371            .iter()
372            .find(|i| i.get("Artist") == Some(&json!({"S": "Foo"})))
373            .expect("item with Artist=Foo");
374
375        // Scalars: whole typed value.
376        assert_eq!(foo.get("str"), Some(&json!({"S": "hello"})));
377        assert_eq!(foo.get("Plays"), Some(&json!({"N": "42"})));
378        assert_eq!(foo.get("bin"), Some(&json!({"B": "aGVsbG8="})));
379        assert_eq!(foo.get("flag"), Some(&json!({"BOOL": true})));
380        assert_eq!(foo.get("nothing"), Some(&json!({"NULL": true})));
381
382        // List (3 mixed elements) and map (2 fields): exact whole-value equality.
383        assert_eq!(
384            foo.get("list"),
385            Some(&json!({"L": [{"S": "nested"}, {"N": "7"}, {"BOOL": false}]}))
386        );
387        assert_eq!(
388            foo.get("map"),
389            Some(&json!({"M": {"genre": {"S": "rock"}, "year": {"N": "1994"}}}))
390        );
391
392        // Sets: assert length AND every member (order is preserved verbatim,
393        // but check membership so a dropped-element regression can't hide).
394        let ss = foo.get("strset").unwrap()["SS"].as_array().unwrap();
395        assert_eq!(ss.len(), 2);
396        assert!(ss.contains(&json!("x")));
397        assert!(ss.contains(&json!("y")));
398
399        let ns = foo.get("numset").unwrap()["NS"].as_array().unwrap();
400        assert_eq!(ns.len(), 2);
401        assert!(ns.contains(&json!("1")));
402        assert!(ns.contains(&json!("2")));
403
404        let bs = foo.get("binset").unwrap()["BS"].as_array().unwrap();
405        assert_eq!(bs.len(), 2);
406        assert!(bs.contains(&json!("YQ==")));
407        assert!(bs.contains(&json!("Yg==")));
408    }
409
410    #[tokio::test]
411    async fn imported_table_persists_through_snapshot_save() {
412        // Regression: a startup bulk-import mutates state directly, so unless it
413        // is written through the DynamoDB snapshot store it is durable only if a
414        // later mutating API call happens to trigger a save. A read-only workload
415        // would lose the imported table on restart. Exercise import -> save ->
416        // reload (the exact path the server now wires) and assert the table
417        // survives.
418        use crate::save_dynamodb_snapshot;
419        use crate::state::DynamoDbSnapshot;
420        use fakecloud_persistence::{DiskSnapshotStore, SnapshotStore};
421        use std::sync::Arc;
422
423        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
424            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
425        ));
426
427        let fixtures = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
428        let describe = read_json(&fixtures.join("describe-table.json")).unwrap();
429        let export_dir = fixtures.join("export");
430        let outcome =
431            import_aws_export(&state, "123456789012", "us-east-1", &export_dir, &describe)
432                .expect("import should succeed");
433        assert!(matches!(outcome, ImportOutcome::Imported { .. }));
434
435        let tmp = tempfile::TempDir::new().unwrap();
436        let store: Arc<dyn SnapshotStore> = Arc::new(DiskSnapshotStore::new(
437            tmp.path().join("dynamodb").join("snapshot.json"),
438        ));
439        let lock = tokio::sync::Mutex::new(());
440        let wrote = save_dynamodb_snapshot(&state, Some(store.clone()), &lock)
441            .await
442            .expect("snapshot save should succeed");
443        assert!(
444            wrote,
445            "a snapshot store is configured, so a save must occur"
446        );
447
448        // Reload exactly as boot does: read the bytes back and deserialize.
449        let bytes = store
450            .load()
451            .expect("load ok")
452            .expect("snapshot bytes on disk");
453        let snapshot: DynamoDbSnapshot =
454            serde_json::from_slice(&bytes).expect("snapshot deserializes");
455        let accounts = snapshot
456            .accounts
457            .expect("v2 multi-account snapshot written");
458        let account = accounts
459            .get("123456789012")
460            .expect("account present after reload");
461        let table = account
462            .tables
463            .get("Music")
464            .expect("imported table survives snapshot round-trip");
465        assert_eq!(table.items.len(), 3, "all imported items are durable");
466    }
467
468    #[test]
469    fn rejects_items_missing_declared_key_attribute() {
470        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
471            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
472        ));
473
474        let export_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/export");
475
476        // KeySchema names a partition key that the exported items don't carry;
477        // the import must fail rather than materialise a key-corrupt table.
478        let describe = json!({
479            "Table": {
480                "TableName": "Music",
481                "KeySchema": [{ "AttributeName": "DoesNotExist", "KeyType": "HASH" }],
482                "AttributeDefinitions": [
483                    { "AttributeName": "DoesNotExist", "AttributeType": "S" }
484                ],
485                "BillingMode": "PAY_PER_REQUEST"
486            }
487        });
488
489        let err = import_aws_export(&state, "123456789012", "us-east-1", &export_dir, &describe)
490            .expect_err("import should fail when items lack the declared key attribute");
491        assert!(
492            err.contains("DoesNotExist"),
493            "error should name the missing key attribute: {err}"
494        );
495    }
496
497    #[test]
498    fn rejects_key_attribute_type_mismatch() {
499        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
500            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
501        ));
502
503        let export_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/export");
504
505        // The exported items carry Artist as a string (`{"S": ...}`), but this
506        // describe-table declares the partition key as numeric (`N`). The import
507        // must fail with a type mismatch, matching the normal write path's
508        // check_key_type, rather than materialise a table a correctly-typed read
509        // could never query.
510        let describe = json!({
511            "Table": {
512                "TableName": "MusicTypeMismatch",
513                "KeySchema": [
514                    { "AttributeName": "Artist", "KeyType": "HASH" },
515                    { "AttributeName": "SongTitle", "KeyType": "RANGE" }
516                ],
517                "AttributeDefinitions": [
518                    { "AttributeName": "Artist", "AttributeType": "N" },
519                    { "AttributeName": "SongTitle", "AttributeType": "S" }
520                ],
521                "BillingMode": "PAY_PER_REQUEST"
522            }
523        });
524
525        let err = import_aws_export(&state, "123456789012", "us-east-1", &export_dir, &describe)
526            .expect_err("import should fail when a key attribute has the wrong type");
527        assert!(
528            err.contains("Artist") && err.contains("type S") && err.contains("declares N"),
529            "error should describe the key type mismatch: {err}"
530        );
531        // Nothing partial landed in state.
532        let mut guard = state.write();
533        assert!(!guard
534            .get_or_create("123456789012")
535            .tables
536            .contains_key("MusicTypeMismatch"));
537    }
538
539    #[test]
540    fn import_is_idempotent_when_table_already_exists() {
541        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
542            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
543        ));
544
545        let fixtures = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
546        let describe = read_json(&fixtures.join("describe-table.json")).unwrap();
547        let export_dir = fixtures.join("export");
548
549        // First import materialises the table (3 items).
550        let first = import_aws_export(&state, "123456789012", "us-east-1", &export_dir, &describe)
551            .expect("first import should succeed");
552        assert_eq!(
553            first,
554            ImportOutcome::Imported {
555                table: "Music".to_string(),
556                items: 3,
557            }
558        );
559
560        // Replace the loaded data with a single sentinel row, standing in for a
561        // table restored from a persisted snapshot that has since diverged.
562        {
563            let mut guard = state.write();
564            let table = guard
565                .get_or_create("123456789012")
566                .tables
567                .get_mut("Music")
568                .expect("table in state");
569            let mut sentinel: Item = HashMap::new();
570            sentinel.insert("Artist".to_string(), json!({ "S": "SENTINEL" }));
571            sentinel.insert("SongTitle".to_string(), json!({ "S": "only" }));
572            table.items = vec![sentinel];
573            table.recalculate_stats();
574        }
575
576        // Re-running with the same flags must skip (not error, not overwrite).
577        let second = import_aws_export(&state, "123456789012", "us-east-1", &export_dir, &describe)
578            .expect("re-import should be a no-op, not an error");
579        assert_eq!(
580            second,
581            ImportOutcome::SkippedExisting {
582                table: "Music".to_string(),
583            }
584        );
585
586        // Existing data was left untouched: still the single sentinel row, not
587        // re-populated with the 3 export items.
588        let mut guard = state.write();
589        let table = guard
590            .get_or_create("123456789012")
591            .tables
592            .get("Music")
593            .expect("table still in state");
594        assert_eq!(table.items.len(), 1);
595        assert_eq!(
596            table.items[0].get("Artist"),
597            Some(&json!({ "S": "SENTINEL" }))
598        );
599    }
600
601    #[test]
602    fn rejects_manifest_item_count_mismatch() {
603        use flate2::write::GzEncoder;
604        use flate2::Compression;
605        use std::io::Write as _;
606
607        // Build a throwaway export whose manifest-summary over-declares the item
608        // count (99) relative to the single row actually present, simulating a
609        // truncated/corrupt export.
610        let dir = std::env::temp_dir().join(format!("fc-ddb-import-{}", uuid::Uuid::new_v4()));
611        std::fs::create_dir_all(dir.join("data")).unwrap();
612        std::fs::write(
613            dir.join("manifest-summary.json"),
614            r#"{"version":"2020-06-30","exportFormat":"DYNAMODB_JSON","itemCount":99}"#,
615        )
616        .unwrap();
617        // Per-file entry carries no itemCount, so only the summary-level check
618        // can fire here.
619        std::fs::write(
620            dir.join("manifest-files.json"),
621            r#"{"dataFileS3Key":"AWSDynamoDB/01700000000000-abcd/data/0001.json.gz"}"#,
622        )
623        .unwrap();
624        let data_file = std::fs::File::create(dir.join("data/0001.json.gz")).unwrap();
625        let mut enc = GzEncoder::new(data_file, Compression::default());
626        enc.write_all(b"{\"Item\":{\"Artist\":{\"S\":\"A\"},\"SongTitle\":{\"S\":\"B\"}}}\n")
627            .unwrap();
628        enc.finish().unwrap();
629
630        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
631            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
632        ));
633        let describe = json!({
634            "Table": {
635                "TableName": "MusicTruncated",
636                "KeySchema": [
637                    { "AttributeName": "Artist", "KeyType": "HASH" },
638                    { "AttributeName": "SongTitle", "KeyType": "RANGE" }
639                ],
640                "AttributeDefinitions": [
641                    { "AttributeName": "Artist", "AttributeType": "S" },
642                    { "AttributeName": "SongTitle", "AttributeType": "S" }
643                ],
644                "BillingMode": "PAY_PER_REQUEST"
645            }
646        });
647
648        let err = import_aws_export(&state, "123456789012", "us-east-1", &dir, &describe)
649            .expect_err("import should fail when manifest itemCount disagrees with the data");
650        assert!(
651            err.contains("99") && err.contains("1 items"),
652            "error should report the declared vs actual counts: {err}"
653        );
654
655        let _ = std::fs::remove_dir_all(&dir);
656    }
657
658    #[test]
659    fn rejects_per_file_item_count_mismatch() {
660        use flate2::write::GzEncoder;
661        use flate2::Compression;
662        use std::io::Write as _;
663
664        // manifest-summary agrees with the data (1 item), but the per-file
665        // manifest entry over-declares itemCount (5). The per-file integrity
666        // check must catch it.
667        let dir = std::env::temp_dir().join(format!("fc-ddb-import-{}", uuid::Uuid::new_v4()));
668        std::fs::create_dir_all(dir.join("data")).unwrap();
669        std::fs::write(
670            dir.join("manifest-summary.json"),
671            r#"{"version":"2020-06-30","exportFormat":"DYNAMODB_JSON","itemCount":1}"#,
672        )
673        .unwrap();
674        std::fs::write(
675            dir.join("manifest-files.json"),
676            r#"{"itemCount":5,"dataFileS3Key":"AWSDynamoDB/01700000000000-abcd/data/0001.json.gz"}"#,
677        )
678        .unwrap();
679        let data_file = std::fs::File::create(dir.join("data/0001.json.gz")).unwrap();
680        let mut enc = GzEncoder::new(data_file, Compression::default());
681        enc.write_all(b"{\"Item\":{\"Artist\":{\"S\":\"A\"},\"SongTitle\":{\"S\":\"B\"}}}\n")
682            .unwrap();
683        enc.finish().unwrap();
684
685        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
686            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
687        ));
688        let describe = json!({
689            "Table": {
690                "TableName": "MusicPerFile",
691                "KeySchema": [
692                    { "AttributeName": "Artist", "KeyType": "HASH" },
693                    { "AttributeName": "SongTitle", "KeyType": "RANGE" }
694                ],
695                "AttributeDefinitions": [
696                    { "AttributeName": "Artist", "AttributeType": "S" },
697                    { "AttributeName": "SongTitle", "AttributeType": "S" }
698                ],
699                "BillingMode": "PAY_PER_REQUEST"
700            }
701        });
702
703        let err = import_aws_export(&state, "123456789012", "us-east-1", &dir, &describe)
704            .expect_err("import should fail when a data file's itemCount is wrong");
705        assert!(
706            err.contains("0001.json.gz") && err.contains('5'),
707            "error should name the data file and declared count: {err}"
708        );
709
710        let _ = std::fs::remove_dir_all(&dir);
711    }
712
713    /// Writes a minimal self-contained single-item table export under
714    /// `root/<subdir_name>`, as `import_aws_exports_dir` expects.
715    fn write_table_subdir(root: &Path, subdir_name: &str, table_name: &str) {
716        use flate2::write::GzEncoder;
717        use flate2::Compression;
718        use std::io::Write as _;
719
720        let dir = root.join(subdir_name);
721        std::fs::create_dir_all(dir.join("data")).unwrap();
722        std::fs::write(
723            dir.join("describe-table.json"),
724            json!({
725                "Table": {
726                    "TableName": table_name,
727                    "KeySchema": [{ "AttributeName": "Id", "KeyType": "HASH" }],
728                    "AttributeDefinitions": [{ "AttributeName": "Id", "AttributeType": "S" }],
729                    "BillingMode": "PAY_PER_REQUEST"
730                }
731            })
732            .to_string(),
733        )
734        .unwrap();
735        std::fs::write(
736            dir.join("manifest-summary.json"),
737            r#"{"version":"2020-06-30","exportFormat":"DYNAMODB_JSON","itemCount":1}"#,
738        )
739        .unwrap();
740        std::fs::write(
741            dir.join("manifest-files.json"),
742            r#"{"itemCount":1,"dataFileS3Key":"AWSDynamoDB/01700000000000-abcd/data/0001.json.gz"}"#,
743        )
744        .unwrap();
745        let data_file = std::fs::File::create(dir.join("data/0001.json.gz")).unwrap();
746        let mut enc = GzEncoder::new(data_file, Compression::default());
747        enc.write_all(format!(r#"{{"Item":{{"Id":{{"S":"{table_name}-row"}}}}}}"#).as_bytes())
748            .unwrap();
749        enc.write_all(b"\n").unwrap();
750        enc.finish().unwrap();
751    }
752
753    #[test]
754    fn imports_multiple_tables_from_root_dir() {
755        let root = std::env::temp_dir().join(format!("fc-ddb-multi-{}", uuid::Uuid::new_v4()));
756        write_table_subdir(&root, "a-table", "TableA");
757        write_table_subdir(&root, "b-table", "TableB");
758
759        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
760            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
761        ));
762
763        let outcomes = import_aws_exports_dir(&state, "123456789012", "us-east-1", &root)
764            .expect("multi-table import should succeed");
765        assert_eq!(outcomes.len(), 2);
766        assert!(outcomes
767            .iter()
768            .all(|o| matches!(o, ImportOutcome::Imported { items: 1, .. })));
769
770        let mut guard = state.write();
771        let account = guard.get_or_create("123456789012");
772        assert!(account.tables.contains_key("TableA"));
773        assert!(account.tables.contains_key("TableB"));
774        drop(guard);
775
776        let _ = std::fs::remove_dir_all(&root);
777    }
778
779    #[test]
780    fn multi_table_import_fails_when_root_has_no_subdirectories() {
781        let root =
782            std::env::temp_dir().join(format!("fc-ddb-multi-empty-{}", uuid::Uuid::new_v4()));
783        std::fs::create_dir_all(&root).unwrap();
784
785        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
786            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
787        ));
788        let err = import_aws_exports_dir(&state, "123456789012", "us-east-1", &root)
789            .expect_err("empty root should fail");
790        assert!(err.contains("no per-table subdirectories"));
791
792        let _ = std::fs::remove_dir_all(&root);
793    }
794
795    #[test]
796    fn multi_table_import_fails_on_subdir_missing_describe_table() {
797        let root = std::env::temp_dir().join(format!("fc-ddb-multi-bad-{}", uuid::Uuid::new_v4()));
798        std::fs::create_dir_all(root.join("bad-table")).unwrap();
799
800        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
801            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
802        ));
803        let err = import_aws_exports_dir(&state, "123456789012", "us-east-1", &root)
804            .expect_err("subdir without describe-table.json should fail");
805        assert!(
806            err.contains("bad-table"),
807            "error should name the offending subdirectory: {err}"
808        );
809
810        let _ = std::fs::remove_dir_all(&root);
811    }
812
813    #[test]
814    fn multi_table_import_stops_at_first_failure_leaving_earlier_tables_committed() {
815        let root =
816            std::env::temp_dir().join(format!("fc-ddb-multi-partial-{}", uuid::Uuid::new_v4()));
817        write_table_subdir(&root, "a-table", "TableA");
818        std::fs::create_dir_all(root.join("b-table")).unwrap(); // no describe-table.json
819
820        let state: SharedDynamoDbState = std::sync::Arc::new(parking_lot::RwLock::new(
821            fakecloud_core::multi_account::MultiAccountState::new("123456789012", "us-east-1", ""),
822        ));
823        let err = import_aws_exports_dir(&state, "123456789012", "us-east-1", &root)
824            .expect_err("second table should fail");
825        assert!(err.contains("b-table"));
826
827        // "a-table" sorts before "b-table", so it was already imported and
828        // committed to state before the second subdirectory failed.
829        let mut guard = state.write();
830        assert!(guard
831            .get_or_create("123456789012")
832            .tables
833            .contains_key("TableA"));
834        drop(guard);
835
836        let _ = std::fs::remove_dir_all(&root);
837    }
838}