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;
5use async_trait::async_trait;
6use base64::{prelude::BASE64_STANDARD, Engine};
7use chrono::Utc;
8use std::collections::HashMap;
9use std::fmt::{Debug, Formatter};
10
11use super::{validate_key, validate_value};
12
13/// AWS DynamoDB implementation of the KV trait.
14///
15/// Credential refresh is handled automatically by the underlying `AwsCredentialProvider`
16/// inside `DynamoDbClient`.
17pub struct AwsDynamodbKv {
18    client: DynamoDbClient,
19    table_name: String,
20}
21
22impl Debug for AwsDynamodbKv {
23    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
24        f.debug_struct("AwsDynamodbKv")
25            .field("table_name", &self.table_name)
26            .finish()
27    }
28}
29
30impl AwsDynamodbKv {
31    pub fn new(table_name: String, client: DynamoDbClient) -> Self {
32        Self { client, table_name }
33    }
34
35    /// Creates a hash bucket for load distribution
36    fn hash_bucket(&self, key: &str) -> String {
37        use std::collections::hash_map::DefaultHasher;
38        use std::hash::{Hash, Hasher};
39
40        let mut hasher = DefaultHasher::new();
41        key.hash(&mut hasher);
42        let bucket_id = hasher.finish() % 16; // 16 buckets for load distribution
43        format!("bucket_{}", bucket_id)
44    }
45
46    /// Checks if an item has expired based on TTL
47    fn is_expired(&self, ttl_epoch: Option<i64>) -> bool {
48        if let Some(ttl_timestamp) = ttl_epoch {
49            let now = Utc::now().timestamp();
50            now >= ttl_timestamp
51        } else {
52            false
53        }
54    }
55}
56
57impl Binding for AwsDynamodbKv {}
58
59#[async_trait]
60impl Kv for AwsDynamodbKv {
61    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
62        validate_key(key)?;
63
64        let bucket = self.hash_bucket(key);
65        let mut primary_key = HashMap::new();
66        primary_key.insert("pk".to_string(), AttributeValue::s(bucket));
67        primary_key.insert("sk".to_string(), AttributeValue::s(key.to_string()));
68
69        let request = GetItemRequest::builder()
70            .table_name(self.table_name.clone())
71            .key(primary_key)
72            // `Kv::put` followed by `Kv::get` must observe the write. DynamoDB
73            // GetItem is eventually consistent by default, which can make a
74            // freshly stored command payload appear missing during immediate
75            // push dispatch.
76            .consistent_read(true)
77            .build();
78
79        let response = self.client.get_item(request).await.map_err(|e| {
80            map_cloud_client_error(
81                e,
82                format!("Failed to get item with key '{}'", key),
83                Some(key.to_string()),
84            )
85        })?;
86
87        if let Some(item) = response.item {
88            // Check TTL expiry (logical expiry contract)
89            if let Some(ttl_attr) = item.get("ttl") {
90                if let Some(ttl_epoch) = ttl_attr.n.as_ref().and_then(|s| s.parse::<i64>().ok()) {
91                    if self.is_expired(Some(ttl_epoch)) {
92                        return Ok(None); // Logically expired
93                    }
94                }
95            }
96
97            let value = item
98                .get("value")
99                .and_then(|attr| attr.b.as_ref())
100                .and_then(|base64_value| BASE64_STANDARD.decode(base64_value).ok())
101                .ok_or_else(|| {
102                    AlienError::new(ErrorData::CloudPlatformError {
103                        message: format!("Missing or invalid value attribute for key '{}'", key),
104                        resource_id: Some(key.to_string()),
105                    })
106                })?;
107
108            Ok(Some(value))
109        } else {
110            Ok(None)
111        }
112    }
113
114    async fn put(&self, key: &str, value: Vec<u8>, options: Option<PutOptions>) -> Result<bool> {
115        validate_key(key)?;
116        validate_value(&value)?;
117
118        let bucket = self.hash_bucket(key);
119        let options = options.unwrap_or_default();
120
121        let mut item = HashMap::new();
122        item.insert("pk".to_string(), AttributeValue::s(bucket));
123        item.insert("sk".to_string(), AttributeValue::s(key.to_string()));
124        item.insert(
125            "value".to_string(),
126            AttributeValue::b(BASE64_STANDARD.encode(&value)),
127        );
128
129        if let Some(ttl) = options.ttl {
130            let expires_at = (Utc::now() + ttl).timestamp();
131            item.insert("ttl".to_string(), AttributeValue::n(expires_at.to_string()));
132        }
133
134        let request = if options.if_not_exists {
135            // Expired rows count as ABSENT, exactly like the local provider's
136            // atomic takeover: DynamoDB's background TTL sweeper can lag the
137            // logical expiry by hours, and without the `#ttl <= :now` arm a
138            // conditional put (e.g. a command lease takeover after the
139            // previous holder died) would be blocked until the sweeper
140            // physically deletes the row.
141            let mut expression_attribute_names = HashMap::new();
142            expression_attribute_names.insert("#ttl".to_string(), "ttl".to_string());
143            let mut expression_attribute_values = HashMap::new();
144            expression_attribute_values.insert(
145                ":now".to_string(),
146                AttributeValue::n(Utc::now().timestamp().to_string()),
147            );
148            PutItemRequest::builder()
149                .table_name(self.table_name.clone())
150                .item(item)
151                .condition_expression(
152                    "(attribute_not_exists(pk) AND attribute_not_exists(sk))                      OR (attribute_exists(#ttl) AND #ttl <= :now)"
153                        .to_string(),
154                )
155                .expression_attribute_names(expression_attribute_names)
156                .expression_attribute_values(expression_attribute_values)
157                .build()
158        } else {
159            PutItemRequest::builder()
160                .table_name(self.table_name.clone())
161                .item(item)
162                .build()
163        };
164
165        match self.client.put_item(request).await {
166            Ok(_) => Ok(true),
167            Err(e) => {
168                // Check if this is a conditional check failure for if_not_exists
169                if options.if_not_exists {
170                    if let Some(alien_client_core::ErrorData::RemoteResourceConflict { .. }) =
171                        &e.error
172                    {
173                        return Ok(false);
174                    }
175                }
176                Err(map_cloud_client_error(
177                    e,
178                    format!("Failed to put item with key '{}'", key),
179                    Some(key.to_string()),
180                ))
181            }
182        }
183    }
184
185    async fn delete(&self, key: &str) -> Result<()> {
186        validate_key(key)?;
187
188        let bucket = self.hash_bucket(key);
189        let mut primary_key = HashMap::new();
190        primary_key.insert("pk".to_string(), AttributeValue::s(bucket));
191        primary_key.insert("sk".to_string(), AttributeValue::s(key.to_string()));
192
193        let request = DeleteItemRequest::builder()
194            .table_name(self.table_name.clone())
195            .key(primary_key)
196            .build();
197
198        self.client.delete_item(request).await.map_err(|e| {
199            map_cloud_client_error(
200                e,
201                format!("Failed to delete item with key '{}'", key),
202                Some(key.to_string()),
203            )
204        })?;
205
206        Ok(())
207    }
208
209    async fn exists(&self, key: &str) -> Result<bool> {
210        validate_key(key)?;
211
212        let bucket = self.hash_bucket(key);
213        let mut primary_key = HashMap::new();
214        primary_key.insert("pk".to_string(), AttributeValue::s(bucket));
215        primary_key.insert("sk".to_string(), AttributeValue::s(key.to_string()));
216
217        // Use expression attribute names to avoid reserved keyword 'ttl'
218        let mut expression_attribute_names = HashMap::new();
219        expression_attribute_names.insert("#ttl".to_string(), "ttl".to_string());
220
221        let request = GetItemRequest::builder()
222            .table_name(self.table_name.clone())
223            .key(primary_key)
224            .projection_expression("pk, #ttl".to_string()) // Get key and TTL for expiry check
225            .expression_attribute_names(expression_attribute_names)
226            .consistent_read(true)
227            .build();
228
229        let response = self.client.get_item(request).await.map_err(|e| {
230            map_cloud_client_error(
231                e,
232                format!("Failed to check existence of item with key '{}'", key),
233                Some(key.to_string()),
234            )
235        })?;
236
237        if let Some(item) = response.item {
238            // Check TTL expiry (logical expiry contract)
239            if let Some(ttl_attr) = item.get("ttl") {
240                if let Some(ttl_epoch) = ttl_attr.n.as_ref().and_then(|s| s.parse::<i64>().ok()) {
241                    if self.is_expired(Some(ttl_epoch)) {
242                        return Ok(false); // Logically expired
243                    }
244                }
245            }
246            Ok(true)
247        } else {
248            Ok(false)
249        }
250    }
251
252    async fn scan_prefix(
253        &self,
254        prefix: &str,
255        limit: Option<usize>,
256        _cursor: Option<String>,
257    ) -> Result<ScanResult> {
258        validate_key(prefix)?; // Prefix follows same key validation rules
259
260        // For prefix scans with hash-based bucketing, we must query ALL buckets
261        // since items with the same prefix can be distributed across different buckets
262        let mut all_items = Vec::new();
263        let mut total_fetched = 0;
264        let limit = limit.unwrap_or(1000);
265
266        // For simplicity, we'll query all 16 buckets sequentially
267        // In production, this could be parallelized for better performance
268        for bucket_id in 0..16 {
269            if total_fetched >= limit {
270                break;
271            }
272
273            let bucket = format!("bucket_{}", bucket_id);
274            let mut expression_attribute_values = HashMap::new();
275            expression_attribute_values.insert(":bucket".to_string(), AttributeValue::s(bucket));
276            expression_attribute_values
277                .insert(":prefix".to_string(), AttributeValue::s(prefix.to_string()));
278
279            // Build request for this bucket
280            let request = QueryRequest::builder()
281                .table_name(self.table_name.clone())
282                .key_condition_expression("pk = :bucket AND begins_with(sk, :prefix)".to_string())
283                .expression_attribute_values(expression_attribute_values)
284                .limit((limit - total_fetched) as i32)
285                .build();
286
287            let response = self.client.query(request).await.map_err(|e| {
288                map_cloud_client_error(
289                    e,
290                    format!("Failed to scan prefix '{}' in bucket {}", prefix, bucket_id),
291                    Some(prefix.to_string()),
292                )
293            })?;
294
295            // Process items from this bucket
296            for item in response.items {
297                if total_fetched >= limit {
298                    break;
299                }
300
301                // Check TTL expiry
302                if let Some(ttl_attr) = item.get("ttl") {
303                    if let Some(ttl_epoch) = ttl_attr.n.as_ref().and_then(|s| s.parse::<i64>().ok())
304                    {
305                        if self.is_expired(Some(ttl_epoch)) {
306                            continue; // Skip expired items
307                        }
308                    }
309                }
310
311                if let (Some(key_attr), Some(value_attr)) = (item.get("sk"), item.get("value")) {
312                    if let (Some(key), Some(base64_value)) =
313                        (key_attr.s.as_ref(), value_attr.b.as_ref())
314                    {
315                        if let Ok(value) = BASE64_STANDARD.decode(base64_value) {
316                            all_items.push((key.clone(), value));
317                            total_fetched += 1;
318                        }
319                    }
320                }
321            }
322        }
323
324        // For simplicity, we're not implementing cursor-based pagination across buckets
325        // In production, this would require more complex cursor state management
326        Ok(ScanResult {
327            items: all_items,
328            next_cursor: None,
329        })
330    }
331}
alien_bindings/providers/kv/aws_dynamodb.rs

alien_bindings/providers/kv/
aws_dynamodb.rs
