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#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct TableConfig {
20 #[serde(default = "default_aws_region")]
22 pub region: String,
23 #[serde(default)]
26 pub endpoint: Option<String>,
27 #[serde(default)]
30 pub delete_on_exit: bool,
31 #[serde(default)]
33 pub tables: Vec<TableInfo>,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct TableInfo {
41 pub table_name: String,
45 pub pk: TableAttr,
47 #[serde(default)]
49 pub sk: Option<TableAttr>,
50 #[serde(default)]
53 pub attrs: Vec<TableAttr>,
54 #[serde(default)]
56 pub gsis: Vec<TableGsi>,
57 #[serde(default)]
59 pub lsis: Vec<TableLsi>,
60 #[serde(default)]
62 pub throughput: Option<Throughput>,
63 #[serde(default)]
65 pub seed_data_file: Option<String>,
66}
67
68#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct Throughput {
71 pub read: i64,
73 pub write: i64,
75}
76
77#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct TableAttr {
80 pub name: String,
82 #[serde(rename = "type")]
84 pub attr_type: AttrType,
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
89pub enum AttrType {
90 S,
92 N,
94 B,
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct TableGsi {
101 pub name: String,
103 pub pk: TableAttr,
105 #[serde(default)]
107 pub sk: Option<TableAttr>,
108 #[serde(default)]
110 pub attrs: Vec<String>,
111 #[serde(default)]
113 pub throughput: Option<Throughput>,
114}
115
116#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct TableLsi {
119 pub name: String,
121 pub pk: TableAttr,
123 pub sk: TableAttr,
125 #[serde(default)]
127 pub attrs: Vec<String>,
128}
129
130fn 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 let mut attribute_map: HashMap<String, TableAttr> = HashMap::new();
237
238 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 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 for lsi in &config.lsis {
254 attribute_map.insert(lsi.sk.name.clone(), lsi.sk.clone());
255 }
256
257 let final_attrs: Vec<AttributeDefinition> = attribute_map
259 .into_values()
260 .map(AttributeDefinition::from)
261 .collect();
262
263 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 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 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)); 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 builder.build().map_err(DynamoToolsError::AwsSdkConfig)
320 }
321}
322
323impl TableConfig {
324 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 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 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 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}