Skip to main content

rusty_cdk_core/dynamodb/
builder.rs

1use crate::dynamodb::{AttributeDefinition, KeySchema, Table, TableProperties};
2use crate::dynamodb::{OnDemandThroughput, ProvisionedThroughput, TableRef};
3use crate::shared::{DeletionPolicy, Id, UpdateDeletePolicyDTO, UpdateReplacePolicy};
4use crate::stack::{Resource, StackBuilder};
5use crate::wrappers::{NonZeroNumber, StringWithOnlyAlphaNumericsAndUnderscores};
6use std::marker::PhantomData;
7use crate::type_state;
8
9#[derive(Debug, Clone, 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
24#[derive(Debug, Clone)]
25pub enum AttributeType {
26    String,
27    Number,
28    Binary,
29}
30
31impl From<AttributeType> for String {
32    fn from(value: AttributeType) -> Self {
33        match value {
34            AttributeType::String => "S".to_string(),
35            AttributeType::Number => "N".to_string(),
36            AttributeType::Binary => "B".to_string(),
37        }
38    }
39}
40
41pub struct Key {
42    key: String,
43    key_type: AttributeType,
44}
45
46impl Key {
47    pub fn new(key: StringWithOnlyAlphaNumericsAndUnderscores, key_type: AttributeType) -> Self {
48        Self { key: key.0, key_type }
49    }
50}
51
52type_state!(
53    TableBuilderState,
54    StartState,
55    ProvisionedStateStart,
56    ProvisionedStateReadSet,
57    ProvisionedStateWriteSet,
58    PayPerRequestState,
59);
60
61/// Builder for DynamoDB tables.
62///
63/// Supports both pay-per-request and provisioned billing modes. The builder enforces the correct configuration based on the chosen billing mode.
64///
65/// # Example
66///
67/// ```rust
68/// use rusty_cdk_core::stack::StackBuilder;
69/// use rusty_cdk_core::dynamodb::{TableBuilder, Key, AttributeType};
70/// use rusty_cdk_core::wrappers::*;
71/// use rusty_cdk_macros::{string_with_only_alphanumerics_and_underscores, non_zero_number};
72///
73/// let mut stack_builder = StackBuilder::new();
74///
75/// let partition_key = Key::new(
76///     string_with_only_alphanumerics_and_underscores!("id"),
77///     AttributeType::String
78/// );
79/// let sort_key = Key::new(
80///     string_with_only_alphanumerics_and_underscores!("timestamp"),
81///     AttributeType::Number
82/// );
83///
84/// let on_demand_table = TableBuilder::new("on-demand-table", partition_key)
85///     .sort_key(sort_key)
86///     .pay_per_request_billing()
87///     .build(&mut stack_builder);
88/// ```
89pub struct TableBuilder<T: TableBuilderState> {
90    state: PhantomData<T>,
91    id: Id,
92    table_name: Option<String>,
93    partition_key: Option<Key>,
94    sort_key: Option<Key>,
95    billing_mode: Option<BillingMode>,
96    read_capacity: Option<u32>,
97    write_capacity: Option<u32>,
98    max_read_capacity: Option<u32>,
99    max_write_capacity: Option<u32>,
100    deletion_policy: Option<String>,
101    update_replace_policy: Option<String>,
102}
103
104impl TableBuilder<StartState> {
105    /// Creates a new DynamoDB table builder.
106    ///
107    /// # Arguments
108    /// * `id` - Unique identifier for the table
109    /// * `key` - Partition key definition for the table
110    pub fn new(id: &str, key: Key) -> Self {
111        TableBuilder {
112            state: Default::default(),
113            id: Id(id.to_string()),
114            table_name: None,
115            partition_key: Some(key),
116            sort_key: None,
117            billing_mode: None,
118            read_capacity: None,
119            write_capacity: None,
120            max_read_capacity: None,
121            max_write_capacity: None,
122            deletion_policy: None,
123            update_replace_policy: None,
124        }
125    }
126}
127
128impl<T: TableBuilderState> TableBuilder<T> {
129    pub fn sort_key(self, key: Key) -> Self {
130        Self {
131            sort_key: Some(key),
132            ..self
133        }
134    }
135
136    pub fn table_name(self, name: StringWithOnlyAlphaNumericsAndUnderscores) -> Self {
137        Self {
138            table_name: Some(name.0),
139            ..self
140        }
141    }
142
143    pub fn update_replace_and_deletion_policy(self, update_replace_policy: UpdateReplacePolicy, deletion_policy: DeletionPolicy) -> Self {
144        Self {
145            deletion_policy: Some(deletion_policy.into()),
146            update_replace_policy: Some(update_replace_policy.into()),
147            ..self
148        }
149    }
150
151    /// Configures the table to use pay-per-request billing.
152    ///
153    /// With this mode, you pay per request and can optionally set max throughput limits.
154    pub fn pay_per_request_billing(self) -> TableBuilder<PayPerRequestState> {
155        TableBuilder {
156            billing_mode: Some(BillingMode::PayPerRequest),
157            state: Default::default(),
158            id: self.id,
159            table_name: self.table_name,
160            partition_key: self.partition_key,
161            sort_key: self.sort_key,
162            max_read_capacity: self.max_read_capacity,
163            max_write_capacity: self.max_write_capacity,
164            deletion_policy: self.deletion_policy,
165            update_replace_policy: self.update_replace_policy,
166            read_capacity: None,
167            write_capacity: None,
168        }
169    }
170
171    /// Configures the table to use provisioned billing.
172    ///
173    /// With this mode, you must specify read and write capacity units.
174    pub fn provisioned_billing(self) -> TableBuilder<ProvisionedStateStart> {
175        TableBuilder {
176            billing_mode: Some(BillingMode::Provisioned),
177            state: Default::default(),
178            id: self.id,
179            table_name: self.table_name,
180            partition_key: self.partition_key,
181            sort_key: self.sort_key,
182            read_capacity: self.read_capacity,
183            write_capacity: self.write_capacity,
184            deletion_policy: self.deletion_policy,
185            update_replace_policy: self.update_replace_policy,
186            max_read_capacity: None,
187            max_write_capacity: None,
188        }
189    }
190
191    fn build_internal(self, stack_builder: &mut StackBuilder) -> TableRef {
192        let Key { key, key_type } = self.partition_key.unwrap();
193        let mut key_schema = vec![KeySchema {
194            attribute_name: key.clone(),
195            key_type: "HASH".to_string(),
196        }];
197        let mut key_attributes = vec![AttributeDefinition {
198            attribute_name: key,
199            attribute_type: key_type.into(),
200        }];
201
202        if let Some(Key { key, key_type }) = self.sort_key {
203            let sort_key = KeySchema {
204                attribute_name: key.clone(),
205                key_type: "RANGE".to_string(),
206            };
207            let sort_key_attributes = AttributeDefinition {
208                attribute_name: key,
209                attribute_type: key_type.into(),
210            };
211            key_schema.push(sort_key);
212            key_attributes.push(sort_key_attributes);
213        }
214
215        let billing_mode = self
216            .billing_mode
217            .expect("billing mode should be set, as this is enforced by the builder");
218
219        let provisioned_throughput = if billing_mode == BillingMode::Provisioned {
220            Some(ProvisionedThroughput {
221                read_capacity: self
222                    .read_capacity
223                    .expect("for provisioned billing mode, read capacity should be set"),
224                write_capacity: self
225                    .write_capacity
226                    .expect("for provisioned billing mode, write capacity should be set"),
227            })
228        } else {
229            None
230        };
231
232        let on_demand_throughput = if billing_mode == BillingMode::PayPerRequest {
233            Some(OnDemandThroughput {
234                max_read_capacity: self.max_read_capacity,
235                max_write_capacity: self.max_write_capacity,
236            })
237        } else {
238            None
239        };
240
241        let properties = TableProperties {
242            key_schema,
243            attribute_definitions: key_attributes,
244            billing_mode: billing_mode.into(),
245            provisioned_throughput,
246            on_demand_throughput,
247        };
248
249        let resource_id = Resource::generate_id("DynamoDBTable");
250        stack_builder.add_resource(Table {
251            id: self.id,
252            resource_id: resource_id.clone(),
253            r#type: "AWS::DynamoDB::Table".to_string(),
254            properties,
255            update_delete_policy_dto: UpdateDeletePolicyDTO { deletion_policy: self.deletion_policy, update_replace_policy: self.update_replace_policy },
256        });
257
258        TableRef::internal_new(resource_id)
259    }
260}
261
262impl TableBuilder<PayPerRequestState> {
263    pub fn max_read_capacity(self, capacity: NonZeroNumber) -> Self {
264        Self {
265            max_read_capacity: Some(capacity.0),
266            ..self
267        }
268    }
269
270    pub fn max_write_capacity(self, capacity: NonZeroNumber) -> Self {
271        Self {
272            max_write_capacity: Some(capacity.0),
273            ..self
274        }
275    }
276
277    pub fn build(self, stack_builder: &mut StackBuilder) -> TableRef {
278        self.build_internal(stack_builder)
279    }
280}
281
282impl TableBuilder<ProvisionedStateStart> {
283    pub fn read_capacity(self, capacity: NonZeroNumber) -> TableBuilder<ProvisionedStateReadSet> {
284        TableBuilder {
285            read_capacity: Some(capacity.0),
286            state: Default::default(),
287            id: self.id,
288            table_name: self.table_name,
289            partition_key: self.partition_key,
290            sort_key: self.sort_key,
291            billing_mode: self.billing_mode,
292            write_capacity: self.write_capacity,
293            max_read_capacity: self.read_capacity,
294            max_write_capacity: self.max_write_capacity,
295            deletion_policy: self.deletion_policy,
296            update_replace_policy: self.update_replace_policy,
297        }
298    }
299}
300
301impl TableBuilder<ProvisionedStateReadSet> {
302    pub fn write_capacity(self, capacity: NonZeroNumber) -> TableBuilder<ProvisionedStateWriteSet> {
303        TableBuilder {
304            write_capacity: Some(capacity.0),
305            state: Default::default(),
306            id: self.id,
307            table_name: self.table_name,
308            partition_key: self.partition_key,
309            sort_key: self.sort_key,
310            billing_mode: self.billing_mode,
311            read_capacity: self.read_capacity,
312            max_read_capacity: self.max_read_capacity,
313            max_write_capacity: self.max_write_capacity,
314            deletion_policy: self.deletion_policy,
315            update_replace_policy: self.update_replace_policy,
316        }
317    }
318}
319
320impl TableBuilder<ProvisionedStateWriteSet> {
321    pub fn build(self, stack_builder: &mut StackBuilder) -> TableRef {
322        self.build_internal(stack_builder)
323    }
324}