rusty_cdk_core/dynamodb/
builder.rs

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
60/// Builder for DynamoDB tables.
61///
62/// Supports both pay-per-request and provisioned billing modes. The builder enforces the correct configuration based on the chosen billing mode.
63///
64/// # Example
65///
66/// ```rust
67/// use rusty_cdk_core::stack::StackBuilder;
68/// use rusty_cdk_core::dynamodb::{TableBuilder, Key, AttributeType};
69/// use rusty_cdk_core::wrappers::*;
70/// use rusty_cdk_macros::{string_with_only_alphanumerics_and_underscores, non_zero_number};
71///
72/// let mut stack_builder = StackBuilder::new();
73///
74/// let partition_key = Key::new(
75///     string_with_only_alphanumerics_and_underscores!("id"),
76///     AttributeType::String
77/// );
78/// let sort_key = Key::new(
79///     string_with_only_alphanumerics_and_underscores!("timestamp"),
80///     AttributeType::Number
81/// );
82///
83/// let on_demand_table = TableBuilder::new("on-demand-table", partition_key)
84///     .sort_key(sort_key)
85///     .pay_per_request_billing()
86///     .build(&mut stack_builder);
87/// ```
88pub 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    /// Creates a new DynamoDB table builder.
103    ///
104    /// # Arguments
105    /// * `id` - Unique identifier for the table
106    /// * `key` - Partition key definition for the table
107    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    /// Configures the table to use pay-per-request billing.
139    ///
140    /// With this mode, you pay per request and can optionally set max throughput limits.
141    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    /// Configures the table to use provisioned billing.
157    ///
158    /// With this mode, you must specify read and write capacity units.
159    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}