alien_bindings/providers/kv/
aws_dynamodb.rs

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

alien_bindings/providers/kv/
aws_dynamodb.rs
