dynamodb_tester/
config.rs

1use std::{fs::File, io::BufReader, path::Path};
2
3use anyhow::{Error, Result};
4use aws_sdk_dynamodb::{
5    input::CreateTableInput,
6    model::{
7        AttributeDefinition, GlobalSecondaryIndex, KeySchemaElement, KeyType, LocalSecondaryIndex,
8        Projection, ProjectionType, ProvisionedThroughput, ScalarAttributeType,
9    },
10};
11use serde::{Deserialize, Serialize};
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct TableConfig {
15    pub table_name: String,
16    pub pk: TableAttr,
17    #[serde(default)]
18    pub sk: Option<TableAttr>,
19    #[serde(default)]
20    pub attrs: Vec<TableAttr>,
21    #[serde(default)]
22    pub gsis: Vec<TableGsi>,
23    #[serde(default)]
24    pub lsis: Vec<TableLsi>,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct TableAttr {
29    pub name: String,
30    #[serde(rename = "type")]
31    pub attr_type: AttrType,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
35pub enum AttrType {
36    S,
37    N,
38    B,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct TableGsi {
43    pub name: String,
44    pub pk: TableAttr,
45    #[serde(default)]
46    pub sk: Option<TableAttr>,
47    #[serde(default)]
48    pub attrs: Vec<String>,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct TableLsi {
53    pub name: String,
54    // must be the same as the pk of the table
55    pub pk: TableAttr,
56    pub sk: TableAttr,
57    #[serde(default)]
58    pub attrs: Vec<String>,
59}
60
61impl From<AttrType> for ScalarAttributeType {
62    fn from(attr_type: AttrType) -> Self {
63        match attr_type {
64            AttrType::S => ScalarAttributeType::S,
65            AttrType::N => ScalarAttributeType::N,
66            AttrType::B => ScalarAttributeType::B,
67        }
68    }
69}
70
71impl From<TableAttr> for AttributeDefinition {
72    fn from(attr: TableAttr) -> Self {
73        let attr_type = attr.attr_type.into();
74        AttributeDefinition::builder()
75            .attribute_name(attr.name)
76            .attribute_type(attr_type)
77            .build()
78    }
79}
80
81impl TableAttr {
82    fn to_pk(&self) -> KeySchemaElement {
83        KeySchemaElement::builder()
84            .attribute_name(self.name.clone())
85            .key_type(KeyType::Hash)
86            .build()
87    }
88
89    fn to_sk(&self) -> KeySchemaElement {
90        KeySchemaElement::builder()
91            .attribute_name(self.name.clone())
92            .key_type(KeyType::Range)
93            .build()
94    }
95}
96
97impl From<TableGsi> for GlobalSecondaryIndex {
98    fn from(gsi: TableGsi) -> Self {
99        let pk = gsi.pk.to_pk();
100        let sk = gsi.sk.map(|sk| sk.to_sk());
101
102        let key_schema = if let Some(sk) = sk {
103            vec![pk, sk]
104        } else {
105            vec![pk]
106        };
107        let pt = ProvisionedThroughput::builder()
108            .read_capacity_units(5)
109            .write_capacity_units(5)
110            .build();
111        GlobalSecondaryIndex::builder()
112            .set_key_schema(Some(key_schema))
113            .projection(
114                Projection::builder()
115                    .projection_type(ProjectionType::Include)
116                    .set_non_key_attributes(Some(gsi.attrs))
117                    .build(),
118            )
119            .provisioned_throughput(pt)
120            .index_name(gsi.name)
121            .build()
122    }
123}
124
125impl From<TableLsi> for LocalSecondaryIndex {
126    fn from(lsi: TableLsi) -> Self {
127        let pk = lsi.pk.to_pk();
128        let sk = lsi.sk.to_sk();
129        let key_schema = vec![pk, sk];
130        let projection = if lsi.attrs.is_empty() {
131            Projection::builder()
132                .projection_type(ProjectionType::All)
133                .build()
134        } else {
135            Projection::builder()
136                .projection_type(ProjectionType::Include)
137                .set_non_key_attributes(Some(lsi.attrs))
138                .build()
139        };
140        LocalSecondaryIndex::builder()
141            .set_key_schema(Some(key_schema))
142            .projection(projection)
143            .index_name(lsi.name)
144            .build()
145    }
146}
147
148impl TryFrom<TableConfig> for CreateTableInput {
149    type Error = Error;
150    fn try_from(config: TableConfig) -> Result<Self> {
151        let pk = config.pk.to_pk();
152        let sk = config.sk.as_ref().map(|sk| sk.to_sk());
153
154        let key_schema = if let Some(sk) = sk {
155            vec![pk, sk]
156        } else {
157            vec![pk]
158        };
159
160        // add pk and sk to attrs
161        let mut attrs = config.attrs.clone();
162        attrs.push(config.pk);
163        if let Some(sk) = config.sk {
164            attrs.push(sk);
165        }
166        let attrs = attrs.into_iter().map(AttributeDefinition::from).collect();
167
168        let gsis = config
169            .gsis
170            .into_iter()
171            .map(GlobalSecondaryIndex::from)
172            .collect();
173
174        let lsis = config
175            .lsis
176            .into_iter()
177            .map(LocalSecondaryIndex::from)
178            .collect();
179
180        let pt = ProvisionedThroughput::builder()
181            .read_capacity_units(5)
182            .write_capacity_units(5)
183            .build();
184        let input = CreateTableInput::builder()
185            .table_name(config.table_name)
186            .set_key_schema(Some(key_schema))
187            .set_attribute_definitions(Some(attrs))
188            .set_global_secondary_indexes(Some(gsis))
189            .set_local_secondary_indexes(Some(lsis))
190            .provisioned_throughput(pt);
191
192        Ok(input.build()?)
193    }
194}
195
196impl TableConfig {
197    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
198        let file = File::open(path)?;
199        let reader = BufReader::new(file);
200        let config = serde_yaml::from_reader(reader)?;
201        Ok(config)
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    #[test]
210    fn config_could_be_loaded() {
211        let config = TableConfig::load_from_file("fixtures/config.yml").unwrap();
212
213        assert_eq!(config.table_name, "users");
214        assert_eq!(config.pk.name, "pk");
215        assert_eq!(config.pk.attr_type, AttrType::S);
216
217        let input = CreateTableInput::try_from(config).unwrap();
218        assert_eq!(input.attribute_definitions().unwrap().len(), 5);
219        assert_eq!(input.global_secondary_indexes().unwrap().len(), 1);
220        assert_eq!(input.local_secondary_indexes().unwrap().len(), 1);
221    }
222}
dynamodb_tester/config.rs

dynamodb_tester/
config.rs
