alien_bindings/providers/kv/
aws_dynamodb.rs

1use crate::error::{map_cloud_client_error, ErrorData, Result};
2use crate::traits::{Binding, Kv, KvEntry, PutCondition, PutOptions, ScanResult};
3use alien_aws_clients::dynamodb::*;
4use alien_error::{AlienError, Context, IntoAlienError};
5use async_trait::async_trait;
6use base64::{prelude::BASE64_URL_SAFE_NO_PAD, Engine};
7use chrono::Utc;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::fmt::{Debug, Formatter};
11use uuid::Uuid;
12
13use super::{decode_version, encode_version, validate_key, validate_value};
14
15const HASH_BUCKET_COUNT: u8 = 16;
16
17#[derive(Debug, Serialize, Deserialize)]
18#[serde(rename_all = "camelCase")]
19struct CursorState {
20    version: u8,
21    prefix: String,
22    bucket: u8,
23    last_key: Option<String>,
24}
25
26/// AWS DynamoDB implementation of the KV trait.
27///
28/// Credential refresh is handled automatically by the underlying `AwsCredentialProvider`
29/// inside `DynamoDbClient`.
30pub struct AwsDynamodbKv {
31    client: DynamoDbClient,
32    table_name: String,
33}
34
35impl Debug for AwsDynamodbKv {
36    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("AwsDynamodbKv")
38            .field("table_name", &self.table_name)
39            .finish()
40    }
41}
42
43impl AwsDynamodbKv {
44    pub fn new(table_name: String, client: DynamoDbClient) -> Self {
45        Self { client, table_name }
46    }
47
48    /// Creates a hash bucket for load distribution
49    fn hash_bucket(&self, key: &str) -> String {
50        use std::collections::hash_map::DefaultHasher;
51        use std::hash::{Hash, Hasher};
52
53        let mut hasher = DefaultHasher::new();
54        key.hash(&mut hasher);
55        let bucket_id = hasher.finish() % u64::from(HASH_BUCKET_COUNT);
56        format!("bucket_{}", bucket_id)
57    }
58
59    /// Checks if an item has expired based on TTL
60    fn is_expired(&self, ttl_epoch: Option<i64>) -> bool {
61        if let Some(ttl_timestamp) = ttl_epoch {
62            let now = Utc::now().timestamp();
63            now >= ttl_timestamp
64        } else {
65            false
66        }
67    }
68
69    fn primary_key(&self, key: &str) -> HashMap<String, AttributeValue> {
70        HashMap::from([
71            ("pk".to_string(), AttributeValue::s(self.hash_bucket(key))),
72            ("sk".to_string(), AttributeValue::s(key.to_string())),
73        ])
74    }
75
76    async fn load_item(&self, key: &str) -> Result<Option<HashMap<String, AttributeValue>>> {
77        let request = GetItemRequest::builder()
78            .table_name(self.table_name.clone())
79            .key(self.primary_key(key))
80            .consistent_read(true)
81            .build();
82        self.client
83            .get_item(request)
84            .await
85            .map(|response| response.item)
86            .map_err(|error| {
87                map_cloud_client_error(
88                    error,
89                    format!("Failed to get item with key '{}'", key),
90                    Some(key.to_string()),
91                )
92            })
93    }
94
95    fn item_ttl(item: &HashMap<String, AttributeValue>) -> Option<i64> {
96        item.get("ttl")
97            .and_then(|attribute| attribute.n.as_deref())
98            .and_then(|value| value.parse::<i64>().ok())
99    }
100
101    fn item_value(key: &str, item: &HashMap<String, AttributeValue>) -> Result<Vec<u8>> {
102        item.get("value")
103            .and_then(|attribute| attribute.b.as_ref())
104            .and_then(|value| base64::prelude::BASE64_STANDARD.decode(value).ok())
105            .ok_or_else(|| {
106                AlienError::new(ErrorData::CloudPlatformError {
107                    message: format!("Missing or invalid value attribute for key '{}'", key),
108                    resource_id: Some(key.to_string()),
109                })
110            })
111    }
112
113    fn encode_cursor(state: &CursorState) -> Result<String> {
114        let json =
115            serde_json::to_vec(state)
116                .into_alien_error()
117                .context(ErrorData::InvalidInput {
118                    operation_context: "DynamoDB KV scan cursor encoding".to_string(),
119                    details: "Failed to serialize cursor state".to_string(),
120                    field_name: Some("cursor".to_string()),
121                })?;
122        Ok(BASE64_URL_SAFE_NO_PAD.encode(json))
123    }
124
125    fn decode_cursor(prefix: &str, cursor: &str) -> Result<CursorState> {
126        let decoded = BASE64_URL_SAFE_NO_PAD
127            .decode(cursor)
128            .into_alien_error()
129            .context(ErrorData::InvalidInput {
130                operation_context: "DynamoDB KV scan cursor decoding".to_string(),
131                details: "Invalid cursor encoding".to_string(),
132                field_name: Some("cursor".to_string()),
133            })?;
134        let state: CursorState = serde_json::from_slice(&decoded)
135            .into_alien_error()
136            .context(ErrorData::InvalidInput {
137                operation_context: "DynamoDB KV scan cursor decoding".to_string(),
138                details: "Invalid cursor data".to_string(),
139                field_name: Some("cursor".to_string()),
140            })?;
141        if state.version != 1 || state.prefix != prefix || state.bucket >= HASH_BUCKET_COUNT {
142            return Err(AlienError::new(ErrorData::InvalidInput {
143                operation_context: "DynamoDB KV scan cursor validation".to_string(),
144                details: "Cursor does not belong to this prefix scan".to_string(),
145                field_name: Some("cursor".to_string()),
146            }));
147        }
148        if state
149            .last_key
150            .as_ref()
151            .is_some_and(|last_key| !last_key.starts_with(prefix))
152        {
153            return Err(AlienError::new(ErrorData::InvalidInput {
154                operation_context: "DynamoDB KV scan cursor validation".to_string(),
155                details: "Cursor key does not match the scan prefix".to_string(),
156                field_name: Some("cursor".to_string()),
157            }));
158        }
159        Ok(state)
160    }
161}
162
163impl Binding for AwsDynamodbKv {}
164
165#[async_trait]
166impl Kv for AwsDynamodbKv {
167    async fn get(&self, key: &str) -> Result<Option<KvEntry>> {
168        validate_key(key)?;
169
170        let Some(item) = self.load_item(key).await? else {
171            return Ok(None);
172        };
173        let ttl = Self::item_ttl(&item);
174        if self.is_expired(ttl) {
175            return Ok(None);
176        }
177        let backend_version = item
178            .get("version")
179            .and_then(|value| value.s.clone())
180            .ok_or_else(|| {
181                AlienError::new(ErrorData::CloudPlatformError {
182                    message: format!("Missing or invalid version attribute for key '{}'", key),
183                    resource_id: Some(key.to_string()),
184                })
185            })?;
186        Ok(Some(KvEntry {
187            key: key.to_string(),
188            value: Self::item_value(key, &item)?,
189            version: encode_version(
190                key,
191                backend_version,
192                ttl.map(|expires_at| expires_at.saturating_mul(1000)),
193            )?,
194        }))
195    }
196
197    async fn put(&self, key: &str, value: Vec<u8>, options: Option<PutOptions>) -> Result<bool> {
198        validate_key(key)?;
199        validate_value(&value)?;
200
201        let bucket = self.hash_bucket(key);
202        let options = options.unwrap_or_default();
203
204        let mut item = HashMap::new();
205        item.insert("pk".to_string(), AttributeValue::s(bucket));
206        item.insert("sk".to_string(), AttributeValue::s(key.to_string()));
207        item.insert(
208            "value".to_string(),
209            AttributeValue::b(base64::prelude::BASE64_STANDARD.encode(&value)),
210        );
211        item.insert(
212            "version".to_string(),
213            AttributeValue::s(Uuid::new_v4().simple().to_string()),
214        );
215
216        if let Some(ttl) = options.ttl {
217            let expires_at = (Utc::now() + ttl).timestamp();
218            item.insert("ttl".to_string(), AttributeValue::n(expires_at.to_string()));
219        }
220
221        let request = if matches!(options.condition, PutCondition::Absent) {
222            // Expired rows count as ABSENT, exactly like the local provider's
223            // atomic takeover: DynamoDB's background TTL sweeper can lag the
224            // logical expiry by hours, and without the `#ttl <= :now` arm a
225            // conditional put (e.g. a command lease takeover after the
226            // previous holder died) would be blocked until the sweeper
227            // physically deletes the row.
228            let mut expression_attribute_names = HashMap::new();
229            expression_attribute_names.insert("#ttl".to_string(), "ttl".to_string());
230            let mut expression_attribute_values = HashMap::new();
231            expression_attribute_values.insert(
232                ":now".to_string(),
233                AttributeValue::n(Utc::now().timestamp().to_string()),
234            );
235            PutItemRequest::builder()
236                .table_name(self.table_name.clone())
237                .item(item)
238                .condition_expression(
239                    "(attribute_not_exists(pk) AND attribute_not_exists(sk))                      OR (attribute_exists(#ttl) AND #ttl <= :now)"
240                        .to_string(),
241                )
242                .expression_attribute_names(expression_attribute_names)
243                .expression_attribute_values(expression_attribute_values)
244                .build()
245        } else if let PutCondition::Version(ref version) = options.condition {
246            let expected = decode_version(key, version)?;
247            if expected.expired {
248                return Ok(false);
249            }
250            PutItemRequest::builder()
251                .table_name(self.table_name.clone())
252                .item(item)
253                .condition_expression(
254                    "#version = :version AND (attribute_not_exists(#ttl) OR #ttl > :now)"
255                        .to_string(),
256                )
257                .expression_attribute_names(HashMap::from([
258                    ("#version".to_string(), "version".to_string()),
259                    ("#ttl".to_string(), "ttl".to_string()),
260                ]))
261                .expression_attribute_values(HashMap::from([
262                    (
263                        ":version".to_string(),
264                        AttributeValue::s(expected.backend_version),
265                    ),
266                    (
267                        ":now".to_string(),
268                        AttributeValue::n(Utc::now().timestamp().to_string()),
269                    ),
270                ]))
271                .build()
272        } else {
273            PutItemRequest::builder()
274                .table_name(self.table_name.clone())
275                .item(item)
276                .build()
277        };
278
279        match self.client.put_item(request).await {
280            Ok(_) => Ok(true),
281            Err(e) => {
282                if !matches!(options.condition, PutCondition::None) {
283                    if let Some(alien_client_core::ErrorData::RemoteResourceConflict { .. }) =
284                        &e.error
285                    {
286                        return Ok(false);
287                    }
288                }
289                Err(map_cloud_client_error(
290                    e,
291                    format!("Failed to put item with key '{}'", key),
292                    Some(key.to_string()),
293                ))
294            }
295        }
296    }
297
298    async fn delete(&self, key: &str, if_version: Option<&str>) -> Result<bool> {
299        validate_key(key)?;
300
301        let request = if let Some(version) = if_version {
302            let expected = decode_version(key, version)?;
303            if expected.expired {
304                return Ok(false);
305            }
306            DeleteItemRequest::builder()
307                .table_name(self.table_name.clone())
308                .key(self.primary_key(key))
309                .condition_expression(
310                    "#version = :version AND (attribute_not_exists(#ttl) OR #ttl > :now)"
311                        .to_string(),
312                )
313                .expression_attribute_names(HashMap::from([
314                    ("#version".to_string(), "version".to_string()),
315                    ("#ttl".to_string(), "ttl".to_string()),
316                ]))
317                .expression_attribute_values(HashMap::from([
318                    (
319                        ":version".to_string(),
320                        AttributeValue::s(expected.backend_version),
321                    ),
322                    (
323                        ":now".to_string(),
324                        AttributeValue::n(Utc::now().timestamp().to_string()),
325                    ),
326                ]))
327                .build()
328        } else {
329            DeleteItemRequest::builder()
330                .table_name(self.table_name.clone())
331                .key(self.primary_key(key))
332                .build()
333        };
334
335        match self.client.delete_item(request).await {
336            Ok(_) => Ok(true),
337            Err(error)
338                if if_version.is_some()
339                    && matches!(
340                        error.error.as_ref(),
341                        Some(alien_client_core::ErrorData::RemoteResourceConflict { .. })
342                    ) =>
343            {
344                Ok(false)
345            }
346            Err(error) => Err(map_cloud_client_error(
347                error,
348                format!("Failed to delete item with key '{}'", key),
349                Some(key.to_string()),
350            )),
351        }
352    }
353
354    async fn exists(&self, key: &str) -> Result<bool> {
355        validate_key(key)?;
356
357        let bucket = self.hash_bucket(key);
358        let mut primary_key = HashMap::new();
359        primary_key.insert("pk".to_string(), AttributeValue::s(bucket));
360        primary_key.insert("sk".to_string(), AttributeValue::s(key.to_string()));
361
362        // Use expression attribute names to avoid reserved keyword 'ttl'
363        let mut expression_attribute_names = HashMap::new();
364        expression_attribute_names.insert("#ttl".to_string(), "ttl".to_string());
365
366        let request = GetItemRequest::builder()
367            .table_name(self.table_name.clone())
368            .key(primary_key)
369            .projection_expression("pk, #ttl".to_string()) // Get key and TTL for expiry check
370            .expression_attribute_names(expression_attribute_names)
371            .consistent_read(true)
372            .build();
373
374        let response = self.client.get_item(request).await.map_err(|e| {
375            map_cloud_client_error(
376                e,
377                format!("Failed to check existence of item with key '{}'", key),
378                Some(key.to_string()),
379            )
380        })?;
381
382        if let Some(item) = response.item {
383            // Check TTL expiry (logical expiry contract)
384            if let Some(ttl_attr) = item.get("ttl") {
385                if let Some(ttl_epoch) = ttl_attr.n.as_ref().and_then(|s| s.parse::<i64>().ok()) {
386                    if self.is_expired(Some(ttl_epoch)) {
387                        return Ok(false); // Logically expired
388                    }
389                }
390            }
391            Ok(true)
392        } else {
393            Ok(false)
394        }
395    }
396
397    async fn scan_prefix(
398        &self,
399        prefix: &str,
400        limit: Option<usize>,
401        cursor: Option<String>,
402    ) -> Result<ScanResult> {
403        validate_key(prefix)?;
404        let limit = limit.unwrap_or(1000);
405        let initial = cursor
406            .as_deref()
407            .map(|cursor| Self::decode_cursor(prefix, cursor))
408            .transpose()?
409            .unwrap_or(CursorState {
410                version: 1,
411                prefix: prefix.to_string(),
412                bucket: 0,
413                last_key: None,
414            });
415        if limit == 0 {
416            return Ok(ScanResult {
417                items: Vec::new(),
418                next_cursor: cursor,
419            });
420        }
421
422        let mut items = Vec::with_capacity(limit);
423        let mut bucket_id = initial.bucket;
424        let mut last_key = initial.last_key;
425
426        while bucket_id < HASH_BUCKET_COUNT {
427            let bucket = format!("bucket_{}", bucket_id);
428            let mut expression_attribute_values = HashMap::new();
429            expression_attribute_values
430                .insert(":bucket".to_string(), AttributeValue::s(bucket.clone()));
431            expression_attribute_values
432                .insert(":prefix".to_string(), AttributeValue::s(prefix.to_string()));
433
434            let exclusive_start_key = last_key.as_ref().map(|key| {
435                HashMap::from([
436                    ("pk".to_string(), AttributeValue::s(bucket.clone())),
437                    ("sk".to_string(), AttributeValue::s(key.clone())),
438                ])
439            });
440            let request = QueryRequest::builder()
441                .table_name(self.table_name.clone())
442                .key_condition_expression("pk = :bucket AND begins_with(sk, :prefix)".to_string())
443                .expression_attribute_values(expression_attribute_values)
444                .limit(i32::try_from(limit - items.len()).unwrap_or(i32::MAX))
445                .maybe_exclusive_start_key(exclusive_start_key)
446                .build();
447
448            let response = self.client.query(request).await.map_err(|e| {
449                map_cloud_client_error(
450                    e,
451                    format!("Failed to scan prefix '{}' in bucket {}", prefix, bucket_id),
452                    Some(prefix.to_string()),
453                )
454            })?;
455
456            for item in response.items {
457                if let Some(ttl_attr) = item.get("ttl") {
458                    if let Some(ttl_epoch) = ttl_attr.n.as_ref().and_then(|s| s.parse::<i64>().ok())
459                    {
460                        if self.is_expired(Some(ttl_epoch)) {
461                            continue; // Skip expired items
462                        }
463                    }
464                }
465
466                if let (Some(key_attr), Some(value_attr)) = (item.get("sk"), item.get("value")) {
467                    if let (Some(key), Some(base64_value)) =
468                        (key_attr.s.as_ref(), value_attr.b.as_ref())
469                    {
470                        if let Ok(value) = base64::prelude::BASE64_STANDARD.decode(base64_value) {
471                            let ttl = Self::item_ttl(&item);
472                            if let Some(backend_version) =
473                                item.get("version").and_then(|value| value.s.clone())
474                            {
475                                items.push(KvEntry {
476                                    key: key.clone(),
477                                    value,
478                                    version: encode_version(
479                                        key,
480                                        backend_version,
481                                        ttl.map(|expires_at| expires_at.saturating_mul(1000)),
482                                    )?,
483                                });
484                            } else {
485                                return Err(AlienError::new(ErrorData::CloudPlatformError {
486                                    message: format!(
487                                        "Missing or invalid version attribute for key '{}'",
488                                        key
489                                    ),
490                                    resource_id: Some(key.clone()),
491                                }));
492                            }
493                        }
494                    }
495                }
496            }
497
498            let provider_last_key = response
499                .last_evaluated_key
500                .as_ref()
501                .and_then(|key| key.get("sk"))
502                .and_then(|value| value.s.clone());
503
504            if items.len() == limit {
505                let (next_bucket, next_key) = match provider_last_key {
506                    Some(key) => (bucket_id, Some(key)),
507                    None if bucket_id + 1 < HASH_BUCKET_COUNT => (bucket_id + 1, None),
508                    None => {
509                        return Ok(ScanResult {
510                            items,
511                            next_cursor: None,
512                        });
513                    }
514                };
515                return Ok(ScanResult {
516                    items,
517                    next_cursor: Some(Self::encode_cursor(&CursorState {
518                        version: 1,
519                        prefix: prefix.to_string(),
520                        bucket: next_bucket,
521                        last_key: next_key,
522                    })?),
523                });
524            }
525
526            if let Some(key) = provider_last_key {
527                last_key = Some(key);
528            } else {
529                bucket_id += 1;
530                last_key = None;
531            }
532        }
533
534        Ok(ScanResult {
535            items,
536            next_cursor: None,
537        })
538    }
539}
alien_bindings/providers/kv/aws_dynamodb.rs

alien_bindings/providers/kv/
aws_dynamodb.rs
