1use crate::dynamodb::{AttributeDefinition, KeySchema, Table, TableProperties};
2use crate::dynamodb::{OnDemandThroughput, ProvisionedThroughput, TableRef};
3use crate::shared::Id;
4use crate::stack::{Resource, StackBuilder};
5use crate::wrappers::{NonZeroNumber, StringWithOnlyAlphaNumericsAndUnderscores};
6use std::marker::PhantomData;
7use crate::type_state;
8
9#[derive(PartialEq)]
10pub enum BillingMode {
11 PayPerRequest,
12 Provisioned,
13}
14
15impl From<BillingMode> for String {
16 fn from(value: BillingMode) -> Self {
17 match value {
18 BillingMode::PayPerRequest => "PAY_PER_REQUEST".to_string(),
19 BillingMode::Provisioned => "PROVISIONED".to_string(),
20 }
21 }
22}
23
24pub enum AttributeType {
25 String,
26 Number,
27 Binary,
28}
29
30impl From<AttributeType> for String {
31 fn from(value: AttributeType) -> Self {
32 match value {
33 AttributeType::String => "S".to_string(),
34 AttributeType::Number => "N".to_string(),
35 AttributeType::Binary => "B".to_string(),
36 }
37 }
38}
39
40pub struct Key {
41 key: String,
42 key_type: AttributeType,
43}
44
45impl Key {
46 pub fn new(key: StringWithOnlyAlphaNumericsAndUnderscores, key_type: AttributeType) -> Self {
47 Self { key: key.0, key_type }
48 }
49}
50
51type_state!(
52 TableBuilderState,
53 StartState,
54 ProvisionedStateStart,
55 ProvisionedStateReadSet,
56 ProvisionedStateWriteSet,
57 PayPerRequestState,
58);
59
60pub struct TableBuilder<T: TableBuilderState> {
89 state: PhantomData<T>,
90 id: Id,
91 table_name: Option<String>,
92 partition_key: Option<Key>,
93 sort_key: Option<Key>,
94 billing_mode: Option<BillingMode>,
95 read_capacity: Option<u32>,
96 write_capacity: Option<u32>,
97 max_read_capacity: Option<u32>,
98 max_write_capacity: Option<u32>,
99}
100
101impl TableBuilder<StartState> {
102 pub fn new(id: &str, key: Key) -> Self {
108 TableBuilder {
109 state: Default::default(),
110 id: Id(id.to_string()),
111 table_name: None,
112 partition_key: Some(key),
113 sort_key: None,
114 billing_mode: None,
115 read_capacity: None,
116 write_capacity: None,
117 max_read_capacity: None,
118 max_write_capacity: None,
119 }
120 }
121}
122
123impl<T: TableBuilderState> TableBuilder<T> {
124 pub fn sort_key(self, key: Key) -> Self {
125 Self {
126 sort_key: Some(key),
127 ..self
128 }
129 }
130
131 pub fn table_name(self, name: StringWithOnlyAlphaNumericsAndUnderscores) -> Self {
132 Self {
133 table_name: Some(name.0),
134 ..self
135 }
136 }
137
138 pub fn pay_per_request_billing(self) -> TableBuilder<PayPerRequestState> {
142 TableBuilder {
143 billing_mode: Some(BillingMode::PayPerRequest),
144 state: Default::default(),
145 id: self.id,
146 table_name: self.table_name,
147 partition_key: self.partition_key,
148 sort_key: self.sort_key,
149 max_read_capacity: self.max_read_capacity,
150 max_write_capacity: self.max_write_capacity,
151 read_capacity: None,
152 write_capacity: None,
153 }
154 }
155
156 pub fn provisioned_billing(self) -> TableBuilder<ProvisionedStateStart> {
160 TableBuilder {
161 billing_mode: Some(BillingMode::Provisioned),
162 state: Default::default(),
163 id: self.id,
164 table_name: self.table_name,
165 partition_key: self.partition_key,
166 sort_key: self.sort_key,
167 read_capacity: self.read_capacity,
168 write_capacity: self.write_capacity,
169 max_read_capacity: None,
170 max_write_capacity: None,
171 }
172 }
173
174 fn build_internal(self, stack_builder: &mut StackBuilder) -> TableRef {
175 let Key { key, key_type } = self.partition_key.unwrap();
176 let mut key_schema = vec![KeySchema {
177 attribute_name: key.clone(),
178 key_type: "HASH".to_string(),
179 }];
180 let mut key_attributes = vec![AttributeDefinition {
181 attribute_name: key,
182 attribute_type: key_type.into(),
183 }];
184
185 if let Some(Key { key, key_type }) = self.sort_key {
186 let sort_key = KeySchema {
187 attribute_name: key.clone(),
188 key_type: "RANGE".to_string(),
189 };
190 let sort_key_attributes = AttributeDefinition {
191 attribute_name: key,
192 attribute_type: key_type.into(),
193 };
194 key_schema.push(sort_key);
195 key_attributes.push(sort_key_attributes);
196 }
197
198 let billing_mode = self
199 .billing_mode
200 .expect("billing mode should be set, as this is enforced by the builder");
201
202 let provisioned_throughput = if billing_mode == BillingMode::Provisioned {
203 Some(ProvisionedThroughput {
204 read_capacity: self
205 .read_capacity
206 .expect("for provisioned billing mode, read capacity should be set"),
207 write_capacity: self
208 .write_capacity
209 .expect("for provisioned billing mode, write capacity should be set"),
210 })
211 } else {
212 None
213 };
214
215 let on_demand_throughput = if billing_mode == BillingMode::PayPerRequest {
216 Some(OnDemandThroughput {
217 max_read_capacity: self.max_read_capacity,
218 max_write_capacity: self.max_write_capacity,
219 })
220 } else {
221 None
222 };
223
224 let properties = TableProperties {
225 key_schema,
226 attribute_definitions: key_attributes,
227 billing_mode: billing_mode.into(),
228 provisioned_throughput,
229 on_demand_throughput,
230 };
231
232 let resource_id = Resource::generate_id("DynamoDBTable");
233 stack_builder.add_resource(Table {
234 id: self.id,
235 resource_id: resource_id.clone(),
236 r#type: "AWS::DynamoDB::Table".to_string(),
237 properties,
238 });
239
240 TableRef::new(resource_id)
241 }
242}
243
244impl TableBuilder<PayPerRequestState> {
245 pub fn max_read_capacity(self, capacity: NonZeroNumber) -> Self {
246 Self {
247 max_read_capacity: Some(capacity.0),
248 ..self
249 }
250 }
251
252 pub fn max_write_capacity(self, capacity: NonZeroNumber) -> Self {
253 Self {
254 max_write_capacity: Some(capacity.0),
255 ..self
256 }
257 }
258
259 pub fn build(self, stack_builder: &mut StackBuilder) -> TableRef {
260 self.build_internal(stack_builder)
261 }
262}
263
264impl TableBuilder<ProvisionedStateStart> {
265 pub fn read_capacity(self, capacity: NonZeroNumber) -> TableBuilder<ProvisionedStateReadSet> {
266 TableBuilder {
267 read_capacity: Some(capacity.0),
268 state: Default::default(),
269 id: self.id,
270 table_name: self.table_name,
271 partition_key: self.partition_key,
272 sort_key: self.sort_key,
273 billing_mode: self.billing_mode,
274 write_capacity: self.write_capacity,
275 max_read_capacity: self.read_capacity,
276 max_write_capacity: self.max_write_capacity,
277 }
278 }
279}
280
281impl TableBuilder<ProvisionedStateReadSet> {
282 pub fn write_capacity(self, capacity: NonZeroNumber) -> TableBuilder<ProvisionedStateWriteSet> {
283 TableBuilder {
284 write_capacity: Some(capacity.0),
285 state: Default::default(),
286 id: self.id,
287 table_name: self.table_name,
288 partition_key: self.partition_key,
289 sort_key: self.sort_key,
290 billing_mode: self.billing_mode,
291 read_capacity: self.read_capacity,
292 max_read_capacity: self.max_read_capacity,
293 max_write_capacity: self.max_write_capacity,
294 }
295 }
296}
297
298impl TableBuilder<ProvisionedStateWriteSet> {
299 pub fn build(self, stack_builder: &mut StackBuilder) -> TableRef {
300 self.build_internal(stack_builder)
301 }
302}