dynamodb_tools/
config.rs

1use crate::error::{DynamoToolsError, Result};
2use aws_sdk_dynamodb::{
3    operation::create_table::CreateTableInput,
4    types::{
5        AttributeDefinition, BillingMode, GlobalSecondaryIndex, KeySchemaElement, KeyType,
6        LocalSecondaryIndex, Projection, ProjectionType, ProvisionedThroughput,
7        ScalarAttributeType,
8    },
9};
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12use std::{fs::File, io::BufReader, path::Path};
13
14/// Represents the main configuration loaded from a YAML file.
15///
16/// This struct defines the overall settings for connecting to DynamoDB,
17/// including endpoint, region, and definitions for one or more tables.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct TableConfig {
20    /// AWS region to target. Defaults to "us-east-1" if not specified.
21    #[serde(default = "default_aws_region")]
22    pub region: String,
23    /// Optional local endpoint URL (e.g., "http://localhost:8000" for DynamoDB Local).
24    /// If provided, the connector targets this endpoint and uses test credentials.
25    #[serde(default)]
26    pub endpoint: Option<String>,
27    /// If `true` and `endpoint` is set, created tables will be deleted
28    /// when the `DynamodbConnector` is dropped (requires `test_utils` feature).
29    #[serde(default)]
30    pub delete_on_exit: bool,
31    /// A list of table schemas to be managed by the connector.
32    #[serde(default)]
33    pub tables: Vec<TableInfo>,
34}
35
36/// Defines the detailed schema for a single DynamoDB table.
37///
38/// Used within the `tables` list in [`TableConfig`].
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct TableInfo {
41    /// The base name of the table. A unique ID will be appended
42    /// by the connector upon creation (e.g., `my_table` becomes `my_table-unique_id`).
43    /// This base name is used to retrieve the actual table name later.
44    pub table_name: String,
45    /// The primary partition key attribute definition.
46    pub pk: TableAttr,
47    /// Optional primary sort key attribute definition.
48    #[serde(default)]
49    pub sk: Option<TableAttr>,
50    /// Additional attribute definitions beyond the primary keys.
51    /// PK and SK attributes are automatically included, no need to repeat here.
52    #[serde(default)]
53    pub attrs: Vec<TableAttr>,
54    /// Global Secondary Index definitions.
55    #[serde(default)]
56    pub gsis: Vec<TableGsi>,
57    /// Local Secondary Index definitions.
58    #[serde(default)]
59    pub lsis: Vec<TableLsi>,
60    /// Optional provisioned throughput settings. If `None`, uses Pay-Per-Request billing.
61    #[serde(default)]
62    pub throughput: Option<Throughput>,
63    /// Optional path to a JSON file containing an array of items to seed into the table after creation.
64    #[serde(default)]
65    pub seed_data_file: Option<String>,
66}
67
68/// Defines provisioned throughput settings (read/write capacity units).
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct Throughput {
71    /// Read Capacity Units (RCU).
72    pub read: i64,
73    /// Write Capacity Units (WCU).
74    pub write: i64,
75}
76
77/// Defines a single DynamoDB attribute (name and type).
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct TableAttr {
80    /// The name of the attribute.
81    pub name: String,
82    /// The DynamoDB type of the attribute (S, N, B).
83    #[serde(rename = "type")]
84    pub attr_type: AttrType,
85}
86
87/// Represents the possible DynamoDB scalar attribute types.
88#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
89pub enum AttrType {
90    /// String type.
91    S,
92    /// Number type.
93    N,
94    /// Binary type.
95    B,
96}
97
98/// Defines a Global Secondary Index (GSI).
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct TableGsi {
101    /// The name of the GSI.
102    pub name: String,
103    /// The partition key attribute for the GSI.
104    pub pk: TableAttr,
105    /// Optional sort key attribute for the GSI.
106    #[serde(default)]
107    pub sk: Option<TableAttr>,
108    /// Attributes to project into the GSI (only used if projection type is INCLUDE).
109    #[serde(default)]
110    pub attrs: Vec<String>,
111    /// Optional provisioned throughput for the GSI.
112    #[serde(default)]
113    pub throughput: Option<Throughput>,
114}
115
116/// Defines a Local Secondary Index (LSI).
117#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct TableLsi {
119    /// The name of the LSI.
120    pub name: String,
121    /// The partition key attribute (must be the same as the table's PK).
122    pub pk: TableAttr,
123    /// The sort key attribute for the LSI.
124    pub sk: TableAttr,
125    /// Attributes to project into the LSI (only used if projection type is INCLUDE).
126    #[serde(default)]
127    pub attrs: Vec<String>,
128}
129
130// Internal helper function for default region
131fn default_aws_region() -> String {
132    "us-east-1".to_string()
133}
134
135impl From<AttrType> for ScalarAttributeType {
136    fn from(attr_type: AttrType) -> Self {
137        match attr_type {
138            AttrType::S => ScalarAttributeType::S,
139            AttrType::N => ScalarAttributeType::N,
140            AttrType::B => ScalarAttributeType::B,
141        }
142    }
143}
144
145impl From<TableAttr> for AttributeDefinition {
146    fn from(attr: TableAttr) -> Self {
147        let attr_type = attr.attr_type.into();
148        AttributeDefinition::builder()
149            .attribute_name(attr.name)
150            .attribute_type(attr_type)
151            .build()
152            .unwrap()
153    }
154}
155
156impl TableAttr {
157    fn to_pk(&self) -> KeySchemaElement {
158        KeySchemaElement::builder()
159            .attribute_name(self.name.clone())
160            .key_type(KeyType::Hash)
161            .build()
162            .unwrap()
163    }
164
165    fn to_sk(&self) -> KeySchemaElement {
166        KeySchemaElement::builder()
167            .attribute_name(self.name.clone())
168            .key_type(KeyType::Range)
169            .build()
170            .unwrap()
171    }
172}
173
174impl From<TableGsi> for GlobalSecondaryIndex {
175    fn from(gsi: TableGsi) -> Self {
176        let pk = gsi.pk.to_pk();
177        let sk = gsi.sk.map(|sk| sk.to_sk());
178
179        let key_schema = if let Some(sk) = sk {
180            vec![pk, sk]
181        } else {
182            vec![pk]
183        };
184
185        let mut builder = GlobalSecondaryIndex::builder()
186            .set_key_schema(Some(key_schema))
187            .projection(
188                Projection::builder()
189                    .projection_type(ProjectionType::Include)
190                    .set_non_key_attributes(Some(gsi.attrs))
191                    .build(),
192            )
193            .index_name(gsi.name);
194
195        if let Some(throughput) = gsi.throughput {
196            let pt = ProvisionedThroughput::builder()
197                .read_capacity_units(throughput.read)
198                .write_capacity_units(throughput.write)
199                .build()
200                .unwrap();
201            builder = builder.provisioned_throughput(pt);
202        }
203        builder.build().unwrap()
204    }
205}
206
207impl From<TableLsi> for LocalSecondaryIndex {
208    fn from(lsi: TableLsi) -> Self {
209        let pk = lsi.pk.to_pk();
210        let sk = lsi.sk.to_sk();
211        let key_schema = vec![pk, sk];
212        let projection = if lsi.attrs.is_empty() {
213            Projection::builder()
214                .projection_type(ProjectionType::All)
215                .build()
216        } else {
217            Projection::builder()
218                .projection_type(ProjectionType::Include)
219                .set_non_key_attributes(Some(lsi.attrs))
220                .build()
221        };
222        LocalSecondaryIndex::builder()
223            .set_key_schema(Some(key_schema))
224            .projection(projection)
225            .index_name(lsi.name)
226            .build()
227            .unwrap()
228    }
229}
230
231impl TryFrom<TableInfo> for CreateTableInput {
232    type Error = DynamoToolsError;
233
234    fn try_from(config: TableInfo) -> Result<Self> {
235        // Use a HashMap to collect unique attribute definitions by name
236        let mut attribute_map: HashMap<String, TableAttr> = HashMap::new();
237
238        // 1. Add base table keys
239        attribute_map.insert(config.pk.name.clone(), config.pk.clone());
240        if let Some(ref sk) = config.sk {
241            attribute_map.insert(sk.name.clone(), sk.clone());
242        }
243
244        // 2. Add GSI keys
245        for gsi in &config.gsis {
246            attribute_map.insert(gsi.pk.name.clone(), gsi.pk.clone());
247            if let Some(ref sk) = gsi.sk {
248                attribute_map.insert(sk.name.clone(), sk.clone());
249            }
250        }
251
252        // 4. Add LSI keys
253        for lsi in &config.lsis {
254            attribute_map.insert(lsi.sk.name.clone(), lsi.sk.clone());
255        }
256
257        // Convert the unique attributes to AttributeDefinition vector
258        let final_attrs: Vec<AttributeDefinition> = attribute_map
259            .into_values()
260            .map(AttributeDefinition::from)
261            .collect();
262
263        // --- Key Schema (remains the same) ---
264        let pk_schema = config.pk.to_pk();
265        let sk_schema = config.sk.as_ref().map(|sk| sk.to_sk());
266        let key_schema = if let Some(sk) = sk_schema {
267            vec![pk_schema, sk]
268        } else {
269            vec![pk_schema]
270        };
271        // --- End Key Schema ---
272
273        // --- GSI/LSI Conversion (remains the same) ---
274        let gsis: Vec<GlobalSecondaryIndex> = config
275            .gsis
276            .into_iter()
277            .map(GlobalSecondaryIndex::from)
278            .collect();
279        let lsis: Vec<LocalSecondaryIndex> = config
280            .lsis
281            .into_iter()
282            .map(LocalSecondaryIndex::from)
283            .collect();
284        // --- End GSI/LSI ---
285
286        // --- Build CreateTableInput ---
287        let mut builder = CreateTableInput::builder()
288            .table_name(config.table_name)
289            .set_key_schema(Some(key_schema))
290            .set_attribute_definitions(Some(final_attrs)); // Use the final collected attrs
291
292        if !gsis.is_empty() {
293            builder = builder.set_global_secondary_indexes(Some(gsis));
294        }
295        if !lsis.is_empty() {
296            builder = builder.set_local_secondary_indexes(Some(lsis));
297        }
298
299        match config.throughput {
300            Some(throughput) => {
301                let pt = ProvisionedThroughput::builder()
302                    .read_capacity_units(throughput.read)
303                    .write_capacity_units(throughput.write)
304                    .build()
305                    .map_err(|e| {
306                        DynamoToolsError::Internal(format!(
307                            "Failed to build ProvisionedThroughput: {}",
308                            e
309                        ))
310                    })?;
311                builder = builder.provisioned_throughput(pt);
312            }
313            None => {
314                builder = builder.billing_mode(BillingMode::PayPerRequest);
315            }
316        }
317        // --- End Build ---
318
319        builder.build().map_err(DynamoToolsError::AwsSdkConfig)
320    }
321}
322
323impl TableConfig {
324    /// Loads [`TableConfig`] from a YAML file.
325    ///
326    /// Expects a top-level structure with keys like `region`, `endpoint`, `tables` (a list).
327    ///
328    /// # Errors
329    ///
330    /// Returns `Err` if the file cannot be read ([`DynamoToolsError::ConfigRead`])
331    /// or if the YAML content cannot be parsed ([`DynamoToolsError::ConfigParse`]).
332    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
333        let path_ref = path.as_ref();
334        let path_str = path_ref.to_string_lossy().to_string();
335        let file =
336            File::open(path_ref).map_err(|e| DynamoToolsError::ConfigRead(path_str.clone(), e))?;
337        let reader = BufReader::new(file);
338        let config = serde_yaml::from_reader(reader)
339            .map_err(|e| DynamoToolsError::ConfigParse(path_str, e))?;
340        Ok(config)
341    }
342
343    /// Creates a new `TableConfig` programmatically.
344    pub fn new(
345        region: String,
346        endpoint: Option<String>,
347        delete_on_exit: bool,
348        tables: Vec<TableInfo>,
349    ) -> Self {
350        let delete_on_exit = if endpoint.is_some() {
351            delete_on_exit
352        } else {
353            false
354        };
355
356        Self {
357            region,
358            endpoint,
359            delete_on_exit,
360            tables,
361        }
362    }
363}
364
365impl TableInfo {
366    /// Loads [`TableInfo`] directly from a YAML file.
367    ///
368    /// Generally, it's preferred to load the full [`TableConfig`].
369    ///
370    /// # Errors
371    ///
372    /// Returns `Err` if the file cannot be read ([`DynamoToolsError::ConfigRead`])
373    /// or if the YAML content cannot be parsed ([`DynamoToolsError::ConfigParse`]).
374    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
375        let path_ref = path.as_ref();
376        let path_str = path_ref.to_string_lossy().to_string();
377        let file =
378            File::open(path_ref).map_err(|e| DynamoToolsError::ConfigRead(path_str.clone(), e))?;
379        let reader = BufReader::new(file);
380        let info = serde_yaml::from_reader(reader)
381            .map_err(|e| DynamoToolsError::ConfigParse(path_str, e))?;
382        Ok(info)
383    }
384
385    /// Loads [`TableInfo`] directly from a YAML string.
386    ///
387    /// # Errors
388    ///
389    /// Returns `Err` if the YAML string cannot be parsed ([`DynamoToolsError::ConfigParse`]).
390    ///
391    /// # Example
392    ///
393    /// ```rust
394    /// use dynamodb_tools::{TableInfo, AttrType};
395    ///
396    /// let yaml_data = r#"
397    /// table_name: my_simple_table
398    /// pk:
399    ///   name: item_id
400    ///   type: S
401    /// "#;
402    ///
403    /// let table_info = TableInfo::load(yaml_data).unwrap();
404    ///
405    /// assert_eq!(table_info.table_name, "my_simple_table");
406    /// assert_eq!(table_info.pk.name, "item_id");
407    /// assert_eq!(table_info.pk.attr_type, AttrType::S);
408    /// assert!(table_info.sk.is_none());
409    /// ```
410    pub fn load(s: &str) -> Result<Self> {
411        let info = serde_yaml::from_str(s)
412            .map_err(|e| DynamoToolsError::ConfigParse("string input".to_string(), e))?;
413        Ok(info)
414    }
415}
416
417#[cfg(test)]
418mod tests {
419    use super::*;
420
421    #[test]
422    fn config_could_be_loaded() {
423        let config = TableConfig::load_from_file("fixtures/dev.yml").unwrap();
424        assert_eq!(config.region, "us-east-1");
425        assert_eq!(config.endpoint, Some("http://localhost:8000".to_string()));
426        assert!(config.delete_on_exit);
427        assert!(!config.tables.is_empty());
428
429        let info = config.tables[0].clone();
430        assert_eq!(info.table_name, "users");
431        assert_eq!(info.pk.name, "pk");
432        assert_eq!(info.pk.attr_type, AttrType::S);
433        assert!(info.sk.is_some());
434        assert_eq!(info.gsis.len(), 1);
435        assert_eq!(info.lsis.len(), 1);
436    }
437
438    #[test]
439    fn table_info_could_be_loaded() {
440        let info = TableInfo::load_from_file("fixtures/info.yml").unwrap();
441        assert_eq!(info.table_name, "users");
442        assert_eq!(info.pk.name, "pk");
443        assert_eq!(info.pk.attr_type, AttrType::S);
444    }
445}