Skip to main content

dynamodb_facade/item/
key.rs

1use crate::{HasAttribute, IntoAttributeValue};
2
3use super::*;
4
5/// A type-safe wrapper for a DynamoDB key belonging to a specific table.
6///
7/// `Key<TD>` holds only the key attributes (PK, and SK for composite-key tables)
8/// for the table defined by the [`TableDefinition`] `TD`.
9///
10/// A `Key` can only be obtain from a type implementing [`KeyBuilder`] —
11/// typically any [`DynamoDBItem`] — or by extracting it from an [`Item`].
12///
13/// # Examples
14///
15/// Building a key from a [`KeyBuilder`]:
16///
17/// ```
18/// # use dynamodb_facade::test_fixtures::*;
19/// use dynamodb_facade::{DynamoDBItem, Key, KeyBuilder};
20///
21/// fn platform_key<KB: KeyBuilder<PlatformTable>>(key_builder: &KB) -> Key<PlatformTable> {
22///     key_builder.get_key()
23/// }
24///
25/// // User implements KeyBuilder<PlatformTable> because it is a DynamoDBItem<PlatformTable>
26/// let user = sample_user();
27/// let key = platform_key(&user);
28///
29/// let raw = key.into_inner();
30/// assert_eq!(
31///     raw["PK"].as_s().unwrap(),
32///     "USER#user-1"
33/// );
34/// ```
35///
36/// Extracting a key from an item:
37///
38/// ```
39/// # use dynamodb_facade::test_fixtures::*;
40/// use dynamodb_facade::DynamoDBItem;
41///
42/// let item = sample_user().to_item();
43/// let key = item.into_key_only();
44/// let raw = key.into_inner();
45///
46/// assert!(raw.contains_key("PK"));
47/// assert!(raw.contains_key("SK"));
48/// assert!(!raw.contains_key("name")); // payload stripped
49/// ```
50pub struct Key<TD: TableDefinition>(HashMap<String, AttributeValue>, PhantomData<TD>);
51impl<TD: TableDefinition> fmt::Debug for Key<TD> {
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        fmt::Debug::fmt(&self.0, f)
54    }
55}
56impl<TD: TableDefinition> Key<TD> {
57    /// Consumes the key and returns the underlying raw attribute map.
58    ///
59    /// The map contains only the key attributes (PK and SK for composite-key
60    /// tables). Use this when you need to pass the key to raw SDK builders.
61    ///
62    /// # Examples
63    ///
64    /// ```
65    /// # use dynamodb_facade::test_fixtures::*;
66    /// use dynamodb_facade::{DynamoDBItem, Key, KeyBuilder};
67    ///
68    /// let key: Key<PlatformTable> = sample_user().get_key();
69    /// let raw = key.into_inner();
70    /// assert_eq!(raw["PK"].as_s().unwrap(), "USER#user-1");
71    /// assert_eq!(raw["SK"].as_s().unwrap(), "USER");
72    /// ```
73    pub fn into_inner(self) -> HashMap<String, AttributeValue> {
74        self.0
75    }
76}
77
78/// Builds DynamoDB keys from type-safe key IDs.
79///
80/// This trait is automatically implemented for every type that implements
81/// [`DynamoDBItem`]. It provides three methods:
82///
83/// - [`get_key_from_id`](KeyBuilder::get_key_from_id) — construct a [`Key<TD>`]
84///   from a [`KeyId`] without an instance of the type
85/// - [`get_key_id`](KeyBuilder::get_key_id) — extract the logical [`KeyId`] from
86///   an existing instance
87/// - [`get_key`](KeyBuilder::get_key) — convenience: extract the [`KeyId`] from
88///   `self` and immediately build the [`Key<TD>`]
89///
90/// The associated type `KeyId<'id>` is a [`KeyId<PkId, SkId>`] whose concrete
91/// `PkId` and `SkId` types are determined by the item's [`HasAttribute`] impls.
92///
93/// # Examples
94///
95/// Building a key from an existing instance:
96///
97/// ```
98/// # use dynamodb_facade::test_fixtures::*;
99/// use dynamodb_facade::{Key, KeyBuilder};
100///
101/// let user = sample_user();
102/// let key: Key<PlatformTable> = user.get_key();
103///
104/// let raw = key.into_inner();
105/// assert_eq!(raw["PK"].as_s().unwrap(), "USER#user-1");
106/// assert_eq!(raw["SK"].as_s().unwrap(), "USER");
107/// ```
108///
109/// Building a key from a [`KeyId`] without an instance:
110///
111/// ```
112/// # use dynamodb_facade::test_fixtures::*;
113/// use dynamodb_facade::{Key, KeyBuilder, KeyId};
114///
115/// let key: Key<PlatformTable> = User::get_key_from_id(KeyId::pk("user-42"));
116///
117/// let raw = key.into_inner();
118/// assert_eq!(raw["PK"].as_s().unwrap(), "USER#user-42");
119/// assert_eq!(raw["SK"].as_s().unwrap(), "USER");
120/// ```
121pub trait KeyBuilder<TD: TableDefinition> {
122    /// The logical key identifier type for this item, typically a [`KeyId<PkId, SkId>`]
123    /// whose components are derived from the item's [`HasAttribute`] implementations.
124    type KeyId<'id>;
125
126    /// Constructs a [`Key<TD>`] from a [`KeyId`] without requiring an instance of the
127    /// implementing type.
128    ///
129    /// # Examples
130    ///
131    /// ```
132    /// # use dynamodb_facade::test_fixtures::*;
133    /// use dynamodb_facade::{Key, KeyBuilder, KeyId};
134    ///
135    /// let key: Key<PlatformTable> = User::get_key_from_id(KeyId::pk("user-42"));
136    ///
137    /// let raw = key.into_inner();
138    /// assert_eq!(raw["PK"].as_s().unwrap(), "USER#user-42");
139    /// assert_eq!(raw["SK"].as_s().unwrap(), "USER");
140    /// ```
141    fn get_key_from_id(key_id: Self::KeyId<'_>) -> Key<TD>;
142
143    /// Extracts the logical [`KeyId`] from an existing instance of the implementing type.
144    ///
145    /// The returned [`KeyId`] borrows from `self` and can be passed to
146    /// [`get_key_from_id`][KeyBuilder::get_key_from_id] to produce a [`Key<TD>`].
147    ///
148    /// # Examples
149    ///
150    /// ```
151    /// # use dynamodb_facade::test_fixtures::*;
152    /// use dynamodb_facade::{KeyBuilder, KeyId, NoId};
153    ///
154    /// let user = sample_user();
155    /// let key_id: KeyId<&str, NoId> = <User as KeyBuilder<PlatformTable>>::get_key_id(&user);
156    /// ```
157    fn get_key_id(&self) -> Self::KeyId<'_>;
158
159    /// Convenience method that builds a [`Key<TD>`] directly from `self`.
160    ///
161    /// Equivalent to calling [`get_key_id`][KeyBuilder::get_key_id] followed by
162    /// [`get_key_from_id`][KeyBuilder::get_key_from_id]. Prefer this method when you
163    /// have an instance of the item and simply need its DynamoDB primary key.
164    ///
165    /// # Examples
166    ///
167    /// ```
168    /// # use dynamodb_facade::test_fixtures::*;
169    /// use dynamodb_facade::{Key, KeyBuilder};
170    ///
171    /// let user = sample_user();
172    /// let key: Key<PlatformTable> = user.get_key();
173    ///
174    /// let raw = key.into_inner();
175    /// assert_eq!(raw["PK"].as_s().unwrap(), "USER#user-1");
176    /// assert_eq!(raw["SK"].as_s().unwrap(), "USER");
177    /// ```
178    fn get_key(&self) -> Key<TD> {
179        Self::get_key_from_id(self.get_key_id())
180    }
181}
182
183mod key_builder_helper {
184    //! Blanket [`KeyBuilder`] impl and the internal [`KeyBuilderHelper`] trait,
185    //! dispatched by key schema kind.
186    use super::*;
187
188    impl<TD: TableDefinition, T> KeyBuilder<TD> for T
189    where
190        T: KeyBuilderHelper<TD, <TD::KeySchema as KeySchema>::Kind>,
191    {
192        type KeyId<'id> = KeyId<
193            <Self as KeyBuilderHelper<TD, <TD::KeySchema as KeySchema>::Kind>>::PkId<'id>,
194            <Self as KeyBuilderHelper<TD, <TD::KeySchema as KeySchema>::Kind>>::SkId<'id>,
195        >;
196
197        fn get_key_from_id(key_id: Self::KeyId<'_>) -> Key<TD> {
198            Self::get_key_from_id_helper(key_id)
199        }
200
201        fn get_key_id(&self) -> Self::KeyId<'_> {
202            self.get_key_id_helper()
203        }
204    }
205
206    /// Internal helper that builds keys dispatched by [`KeySchemaKind`].
207    pub trait KeyBuilderHelper<TD: TableDefinition, KSK: KeySchemaKind> {
208        type PkId<'pk>;
209        type SkId<'sk>;
210        fn get_key_from_id_helper(key_id: KeyId<Self::PkId<'_>, Self::SkId<'_>>) -> Key<TD>;
211        fn get_key_id_helper(&self) -> KeyId<Self::PkId<'_>, Self::SkId<'_>>;
212    }
213
214    // -- Simple: PK only ------------------------------------------------------
215
216    impl<TD: TableDefinition, T: DynamoDBItem<TD>> KeyBuilderHelper<TD, SimpleKey> for T
217    where
218        TD::KeySchema: SimpleKeySchema,
219        T: HasTableKeyAttributes<TD>,
220        T: HasAttribute<PartitionKeyDefinition<TD>>,
221    {
222        type PkId<'pk> = <Self as HasAttribute<PartitionKeyDefinition<TD>>>::Id<'pk>;
223        type SkId<'sk> = NoId;
224        fn get_key_from_id_helper(key_id: KeyId<Self::PkId<'_>, Self::SkId<'_>>) -> Key<TD> {
225            let pk_value =
226                <Self as HasAttribute<PartitionKeyDefinition<TD>>>::attribute_value(key_id.pk);
227            Key(
228                HashMap::from([(
229                    PartitionKeyDefinition::<TD>::NAME.to_owned(),
230                    pk_value.into_attribute_value(),
231                )]),
232                PhantomData,
233            )
234        }
235
236        fn get_key_id_helper(&self) -> KeyId<Self::PkId<'_>, Self::SkId<'_>> {
237            let pk_id = T::attribute_id(self);
238            KeyId::pk(pk_id)
239        }
240    }
241
242    // -- Composite: PK + SK ---------------------------------------------------
243
244    impl<TD: TableDefinition, T: DynamoDBItem<TD>> KeyBuilderHelper<TD, CompositeKey> for T
245    where
246        TD::KeySchema: CompositeKeySchema,
247        T: HasTableKeyAttributes<TD>,
248        T: HasAttribute<PartitionKeyDefinition<TD>>,
249        T: HasAttribute<SortKeyDefinition<TD>>,
250    {
251        type PkId<'pk> = <Self as HasAttribute<PartitionKeyDefinition<TD>>>::Id<'pk>;
252        type SkId<'sk> = <Self as HasAttribute<SortKeyDefinition<TD>>>::Id<'sk>;
253
254        fn get_key_from_id_helper(key_id: KeyId<Self::PkId<'_>, Self::SkId<'_>>) -> Key<TD> {
255            let pk_value =
256                <Self as HasAttribute<PartitionKeyDefinition<TD>>>::attribute_value(key_id.pk);
257            let sk_value =
258                <Self as HasAttribute<SortKeyDefinition<TD>>>::attribute_value(key_id.sk);
259            Key(
260                HashMap::from([
261                    (
262                        PartitionKeyDefinition::<TD>::NAME.to_owned(),
263                        pk_value.into_attribute_value(),
264                    ),
265                    (
266                        SortKeyDefinition::<TD>::NAME.to_owned(),
267                        sk_value.into_attribute_value(),
268                    ),
269                ]),
270                PhantomData,
271            )
272        }
273
274        fn get_key_id_helper(&self) -> KeyId<Self::PkId<'_>, Self::SkId<'_>> {
275            let pk_id = <Self as HasAttribute<PartitionKeyDefinition<TD>>>::attribute_id(self);
276            let sk_id = <Self as HasAttribute<SortKeyDefinition<TD>>>::attribute_id(self);
277            KeyId::pk(pk_id).sk(sk_id)
278        }
279    }
280}
281
282impl<TD: TableDefinition> From<Key<TD>> for Item<TD> {
283    fn from(value: Key<TD>) -> Self {
284        Item(value.0, PhantomData)
285    }
286}
287
288impl<TD: TableDefinition> Item<TD>
289where
290    Self: KeyItemExtractor<TD, <TD::KeySchema as KeySchema>::Kind>,
291{
292    /// Splits the item into its key and the remaining non-key attributes.
293    ///
294    /// Returns a tuple of `(Key<TD>, HashMap<String, AttributeValue>)` where
295    /// the key contains only the PK (and SK for composite-key tables) and the
296    /// map contains every other attribute that was in the item.
297    ///
298    /// This is meant to allow the direct manipulation of the attribute map
299    /// while enforcing the invariant that an Item always contains valid key
300    /// attributes. Use it in conjunction with [`Item::from_key_and_attributes`]
301    /// to accomplish that.
302    ///
303    /// # Examples
304    ///
305    /// ```
306    /// # use dynamodb_facade::test_fixtures::*;
307    /// use dynamodb_facade::DynamoDBItem;
308    ///
309    /// let item = sample_user().to_item();
310    /// let (key, rest) = item.extract_key();
311    ///
312    /// // Key contains only PK + SK.
313    /// let raw_key = key.into_inner();
314    /// assert!(raw_key.contains_key("PK"));
315    /// assert!(raw_key.contains_key("SK"));
316    ///
317    /// // Remaining map has everything else.
318    /// assert!(rest.contains_key("name"));
319    /// assert!(!rest.contains_key("PK"));
320    /// ```
321    pub fn extract_key(self) -> (Key<TD>, HashMap<String, AttributeValue>) {
322        KeyItemExtractor::extract_key(self)
323    }
324
325    /// Consumes the item and returns only its key, discarding all other attributes.
326    ///
327    /// This is a convenience wrapper around [`extract_key`](Item::extract_key)
328    /// that drops the remaining attribute map.
329    ///
330    /// # Examples
331    ///
332    /// ```
333    /// # use dynamodb_facade::test_fixtures::*;
334    /// use dynamodb_facade::DynamoDBItem;
335    ///
336    /// let item = sample_user().to_item();
337    /// let key = item.into_key_only();
338    ///
339    /// let raw = key.into_inner();
340    /// assert!(raw.contains_key("PK"));
341    /// assert!(raw.contains_key("SK"));
342    /// assert!(!raw.contains_key("name"));
343    /// ```
344    pub fn into_key_only(self) -> Key<TD> {
345        self.extract_key().0
346    }
347}
348
349/// Splits an [`Item<TD>`] into its [`Key<TD>`] and the remaining attribute map.
350pub trait KeyItemExtractor<TD: TableDefinition, KSK: KeySchemaKind> {
351    fn extract_key(self) -> (Key<TD>, HashMap<String, AttributeValue>);
352}
353
354/// Moves a single key attribute from `from_item` into `to_key`.
355fn transfer_keyschema_helper<TD: TableDefinition>(
356    from_item: &mut Item<TD>,
357    to_key: &mut Key<TD>,
358    k: &str,
359) {
360    let v = from_item
361        .0
362        .remove(k)
363        .expect("Item is guaranteed to contain the KeySchema");
364    to_key.0.insert(k.to_owned(), v);
365}
366impl<TD: TableDefinition> KeyItemExtractor<TD, SimpleKey> for Item<TD>
367where
368    TD::KeySchema: SimpleKeySchema,
369{
370    fn extract_key(mut self) -> (Key<TD>, HashMap<String, AttributeValue>) {
371        let mut key = Key(HashMap::with_capacity(1), PhantomData);
372        transfer_keyschema_helper(
373            &mut self,
374            &mut key,
375            <TD::KeySchema as KeySchema>::PartitionKey::NAME,
376        );
377        (key, self.0)
378    }
379}
380impl<TD: TableDefinition> KeyItemExtractor<TD, CompositeKey> for Item<TD>
381where
382    TD::KeySchema: CompositeKeySchema,
383{
384    fn extract_key(mut self) -> (Key<TD>, HashMap<String, AttributeValue>) {
385        let mut key = Key(HashMap::with_capacity(2), PhantomData);
386        transfer_keyschema_helper(
387            &mut self,
388            &mut key,
389            <TD::KeySchema as KeySchema>::PartitionKey::NAME,
390        );
391        transfer_keyschema_helper(
392            &mut self,
393            &mut key,
394            <TD::KeySchema as CompositeKeySchema>::SortKey::NAME,
395        );
396        (key, self.0)
397    }
398}