alien_bindings/providers/kv/
azure_table_storage.rs

1use crate::error::{ErrorData, Result};
2use crate::traits::{Binding, Kv, PutOptions, ScanResult};
3use alien_azure_clients::tables::{
4    AzureTableStorageClient, EntityQueryOptions, TableEntity, TableStorageApi,
5};
6use alien_error::{AlienError, Context, IntoAlienError};
7use async_trait::async_trait;
8use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
9use chrono::{DateTime, Utc};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12use std::collections::HashMap;
13use std::fmt::{Debug, Formatter};
14
15use super::{validate_key, validate_value};
16
17/// Convert a KV operation to a Table Storage entity
18/// This only base64 encodes the raw bytes when creating the properties map, not in memory
19fn create_table_entity(
20    partition_key: String,
21    row_key: String,
22    value: &[u8],
23    expires_at: Option<DateTime<Utc>>,
24) -> TableEntity {
25    let mut properties = HashMap::new();
26
27    // Base64 encode the raw bytes only when storing in the properties map
28    // This keeps the original 32KB limit valid since we're not storing the encoded version in memory
29    properties.insert("Value".to_string(), Value::String(BASE64.encode(value)));
30
31    // Store creation timestamp
32    properties.insert(
33        "CreatedAt".to_string(),
34        Value::String(Utc::now().to_rfc3339()),
35    );
36
37    // Store expiration timestamp if provided
38    if let Some(expiry) = expires_at {
39        properties.insert("ExpiresAt".to_string(), Value::String(expiry.to_rfc3339()));
40    }
41
42    TableEntity {
43        partition_key,
44        row_key,
45        timestamp: None, // Azure will set this
46        properties,
47    }
48}
49
50/// Extract KV value from Table Storage entity
51fn extract_value_from_entity(entity: &TableEntity) -> Result<Vec<u8>> {
52    let value_str = entity
53        .properties
54        .get("Value")
55        .and_then(|v| v.as_str())
56        .ok_or_else(|| {
57            AlienError::new(ErrorData::InvalidInput {
58                operation_context: "Azure Table Storage KV extract value".to_string(),
59                details: "Entity missing Value property or not a string".to_string(),
60                field_name: Some("Value".to_string()),
61            })
62        })?;
63
64    // Decode base64 value
65    BASE64
66        .decode(value_str)
67        .into_alien_error()
68        .context(ErrorData::InvalidInput {
69            operation_context: "Azure Table Storage KV extract value".to_string(),
70            details: "Failed to decode base64 value".to_string(),
71            field_name: Some("Value".to_string()),
72        })
73}
74
75/// Check if entity has expired based on TTL
76fn is_entity_expired(entity: &TableEntity) -> bool {
77    if let Some(expires_at_value) = entity.properties.get("ExpiresAt") {
78        if let Some(expires_at_str) = expires_at_value.as_str() {
79            if let Ok(expires_at) = DateTime::parse_from_rfc3339(expires_at_str) {
80                return Utc::now() > expires_at.with_timezone(&Utc);
81            }
82        }
83    }
84    false
85}
86
87/// Cursor state for pagination across partitions
88#[derive(Serialize, Deserialize)]
89struct CursorState {
90    current_partition: u32,
91    partition_continuation_token: Option<String>, // Azure's NextPartitionKey + NextRowKey combined
92}
93
94/// Azure Table Storage implementation of the KV trait
95pub struct AzureTableStorageKv {
96    client: AzureTableStorageClient,
97    resource_group_name: String,
98    account_name: String,
99    table_name: String,
100    num_partitions: u32,
101}
102
103impl Debug for AzureTableStorageKv {
104    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
105        f.debug_struct("AzureTableStorageKv")
106            .field("resource_group_name", &self.resource_group_name)
107            .field("account_name", &self.account_name)
108            .field("table_name", &self.table_name)
109            .field("num_partitions", &self.num_partitions)
110            .finish()
111    }
112}
113
114impl AzureTableStorageKv {
115    pub fn new(
116        client: AzureTableStorageClient,
117        resource_group_name: String,
118        account_name: String,
119        table_name: String,
120    ) -> Self {
121        Self {
122            client,
123            resource_group_name,
124            account_name,
125            table_name,
126            num_partitions: 16, // 16 partitions for load distribution
127        }
128    }
129
130    /// Creates a hash bucket for load distribution
131    fn hash_bucket(&self, key: &str) -> u32 {
132        use std::collections::hash_map::DefaultHasher;
133        use std::hash::{Hash, Hasher};
134
135        let mut hasher = DefaultHasher::new();
136        key.hash(&mut hasher);
137        hasher.finish() as u32 % self.num_partitions
138    }
139
140    /// Splits key into partition key and row key
141    /// Attempt to take over a logically-expired row after a conditional
142    /// insert conflict. Returns `Ok(true)` only when THIS caller replaced
143    /// the expired row (If-Match on the read ETag); a live row, a lost
144    /// race, or a row deleted in between all resolve to `Ok(false)`.
145    async fn try_take_over_expired(
146        &self,
147        key: &str,
148        entity: &alien_azure_clients::azure::tables::TableEntity,
149    ) -> Result<bool> {
150        use alien_client_core::ErrorData as CloudErrorData;
151
152        let (partition_key, row_key) = self.split_key(key);
153        let existing = match self
154            .client
155            .get_entity(
156                &self.resource_group_name,
157                &self.account_name,
158                &self.table_name,
159                &partition_key,
160                &row_key,
161                None,
162            )
163            .await
164        {
165            Ok(existing) => existing,
166            // Deleted between the conflict and this read: treat as lost —
167            // the caller's next attempt sees a clean slate.
168            Err(e)
169                if matches!(
170                    e.error.as_ref(),
171                    Some(CloudErrorData::RemoteResourceNotFound { .. })
172                ) =>
173            {
174                return Ok(false);
175            }
176            Err(e) => {
177                return Err(crate::error::map_cloud_client_error(
178                    e,
179                    format!("Failed to read existing entity for key '{}'", key),
180                    Some(key.to_string()),
181                ));
182            }
183        };
184
185        if !is_entity_expired(&existing) {
186            return Ok(false);
187        }
188        let Some(etag) = existing
189            .properties
190            .get("odata.etag")
191            .and_then(|value| value.as_str())
192        else {
193            // No ETag on the read entity — cannot replace safely.
194            return Ok(false);
195        };
196
197        match self
198            .client
199            .update_entity(
200                &self.resource_group_name,
201                &self.account_name,
202                &self.table_name,
203                &partition_key,
204                &row_key,
205                entity,
206                Some(alien_azure_clients::azure::tables::ETag::from(etag)),
207            )
208            .await
209        {
210            Ok(_) => Ok(true),
211            // 409/412: someone else replaced or deleted it first.
212            Err(e)
213                if matches!(
214                    e.error.as_ref(),
215                    Some(CloudErrorData::RemoteResourceConflict { .. })
216                        | Some(CloudErrorData::RemoteResourceNotFound { .. })
217                ) =>
218            {
219                Ok(false)
220            }
221            Err(e) => Err(crate::error::map_cloud_client_error(
222                e,
223                format!("Failed to take over expired entity for key '{}'", key),
224                Some(key.to_string()),
225            )),
226        }
227    }
228
229    fn split_key(&self, key: &str) -> (String, String) {
230        // Use hash-based partitioning for load distribution
231        let partition_key = format!("p{}", self.hash_bucket(key));
232        (partition_key, key.to_string())
233    }
234
235    /// Combines partition key and row key back to original key
236    fn combine_key(&self, _partition_key: &str, row_key: &str) -> String {
237        row_key.to_string() // Row key contains the original key
238    }
239
240    /// Encodes cursor state as base64url JSON for safe HTTP transmission
241    fn encode_cursor(&self, state: &CursorState) -> String {
242        let json = serde_json::to_string(state).unwrap();
243        BASE64.encode(json.as_bytes())
244    }
245
246    /// Decodes cursor state from base64url JSON
247    fn decode_cursor(&self, cursor: &str) -> Result<CursorState> {
248        let decoded =
249            BASE64
250                .decode(cursor)
251                .into_alien_error()
252                .context(ErrorData::InvalidInput {
253                    operation_context: "Azure Table Storage KV cursor decoding".to_string(),
254                    details: "Invalid cursor encoding".to_string(),
255                    field_name: Some("cursor".to_string()),
256                })?;
257        let json =
258            String::from_utf8(decoded)
259                .into_alien_error()
260                .context(ErrorData::InvalidInput {
261                    operation_context: "Azure Table Storage KV cursor decoding".to_string(),
262                    details: "Invalid cursor UTF-8".to_string(),
263                    field_name: Some("cursor".to_string()),
264                })?;
265        serde_json::from_str(&json)
266            .into_alien_error()
267            .context(ErrorData::InvalidInput {
268                operation_context: "Azure Table Storage KV cursor decoding".to_string(),
269                details: "Invalid cursor JSON".to_string(),
270                field_name: Some("cursor".to_string()),
271            })
272    }
273}
274
275impl Binding for AzureTableStorageKv {}
276
277#[async_trait]
278impl Kv for AzureTableStorageKv {
279    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
280        validate_key(key)?;
281
282        let (partition_key, row_key) = self.split_key(key);
283
284        match self
285            .client
286            .get_entity(
287                &self.resource_group_name,
288                &self.account_name,
289                &self.table_name,
290                &partition_key,
291                &row_key,
292                None,
293            )
294            .await
295        {
296            Ok(entity) => {
297                // Check if TTL has expired (client-side filtering)
298                if is_entity_expired(&entity) {
299                    return Ok(None); // Expired
300                }
301
302                let value = extract_value_from_entity(&entity)?;
303                Ok(Some(value))
304            }
305            Err(e) => {
306                use alien_client_core::ErrorData as CloudErrorData;
307                match e.error.as_ref() {
308                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(None),
309                    _ => Err(crate::error::map_cloud_client_error(
310                        e,
311                        format!("Failed to get entity for key '{}'", key),
312                        Some(key.to_string()),
313                    )),
314                }
315            }
316        }
317    }
318
319    async fn put(&self, key: &str, value: Vec<u8>, options: Option<PutOptions>) -> Result<bool> {
320        validate_key(key)?;
321        validate_value(&value)?;
322
323        let options = options.unwrap_or_default();
324        let (partition_key, row_key) = self.split_key(key);
325
326        let expires_at = options.ttl.map(|d| Utc::now() + d);
327        let entity =
328            create_table_entity(partition_key.clone(), row_key.clone(), &value, expires_at);
329
330        if options.if_not_exists {
331            match self
332                .client
333                .insert_entity(
334                    &self.resource_group_name,
335                    &self.account_name,
336                    &self.table_name,
337                    &entity,
338                )
339                .await
340            {
341                Ok(_) => Ok(true),
342                Err(e) => {
343                    use alien_client_core::ErrorData as CloudErrorData;
344                    match e.error.as_ref() {
345                        Some(CloudErrorData::RemoteResourceConflict { .. }) => {
346                            // Expired rows count as ABSENT, matching the
347                            // local provider's atomic takeover: Table
348                            // Storage has no server-side TTL, so without
349                            // this an expired row (e.g. a dead command
350                            // lease) blocks conditional puts forever. The
351                            // read + If-Match replace is race-safe: a
352                            // concurrent taker changes the ETag and this
353                            // update loses with a conflict.
354                            self.try_take_over_expired(key, &entity).await
355                        }
356                        _ => Err(crate::error::map_cloud_client_error(
357                            e,
358                            format!("Failed to insert entity for key '{}'", key),
359                            Some(key.to_string()),
360                        )),
361                    }
362                }
363            }
364        } else {
365            // Insert Or Replace (upsert) - matches Azure REST API terminology
366            self.client
367                .insert_or_replace_entity(
368                    &self.resource_group_name,
369                    &self.account_name,
370                    &self.table_name,
371                    &partition_key,
372                    &row_key,
373                    &entity,
374                )
375                .await
376                .map_err(|e| {
377                    crate::error::map_cloud_client_error(
378                        e,
379                        format!("Failed to upsert entity for key '{}'", key),
380                        Some(key.to_string()),
381                    )
382                })?;
383            Ok(true)
384        }
385    }
386
387    async fn delete(&self, key: &str) -> Result<()> {
388        validate_key(key)?;
389
390        let (partition_key, row_key) = self.split_key(key);
391
392        // Delete entity, ignore if not found
393        match self
394            .client
395            .delete_entity(
396                &self.resource_group_name,
397                &self.account_name,
398                &self.table_name,
399                &partition_key,
400                &row_key,
401                None, // No specific ETag constraint
402            )
403            .await
404        {
405            Ok(_) => Ok(()),
406            Err(e) => {
407                use alien_client_core::ErrorData as CloudErrorData;
408                match e.error.as_ref() {
409                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(()), // No error if key doesn't exist
410                    _ => Err(crate::error::map_cloud_client_error(
411                        e,
412                        format!("Failed to delete entity for key '{}'", key),
413                        Some(key.to_string()),
414                    )),
415                }
416            }
417        }
418    }
419
420    async fn exists(&self, key: &str) -> Result<bool> {
421        validate_key(key)?;
422
423        let (partition_key, row_key) = self.split_key(key);
424
425        match self
426            .client
427            .get_entity(
428                &self.resource_group_name,
429                &self.account_name,
430                &self.table_name,
431                &partition_key,
432                &row_key,
433                None,
434            )
435            .await
436        {
437            Ok(entity) => {
438                // Check TTL expiry
439                Ok(!is_entity_expired(&entity))
440            }
441            Err(e) => {
442                use alien_client_core::ErrorData as CloudErrorData;
443                match e.error.as_ref() {
444                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(false),
445                    _ => Err(crate::error::map_cloud_client_error(
446                        e,
447                        format!("Failed to check existence of entity for key '{}'", key),
448                        Some(key.to_string()),
449                    )),
450                }
451            }
452        }
453    }
454
455    async fn scan_prefix(
456        &self,
457        prefix: &str,
458        limit: Option<usize>,
459        cursor: Option<String>,
460    ) -> Result<ScanResult> {
461        validate_key(prefix)?; // Prefix follows same key validation rules
462
463        // For prefix scans with hash-based partitioning, must fan-out across ALL partitions
464        // A RowKey-only filter forces expensive table-wide scans
465
466        // Decode cursor to get partition progress and continuation tokens
467        let cursor_state = cursor.as_ref().map(|c| self.decode_cursor(c)).transpose()?;
468
469        let mut all_items = Vec::new();
470        let mut total_fetched = 0;
471        let limit = limit.unwrap_or(1000);
472
473        // Start from the partition in cursor, or 0 if no cursor
474        let start_partition = cursor_state.as_ref().map_or(0, |cs| cs.current_partition);
475
476        for partition_id in start_partition..self.num_partitions {
477            let partition_key = format!("p{}", partition_id);
478
479            // Build filter with BOTH PartitionKey and RowKey conditions
480            // Use a range query approach that's compatible with Azure Table Storage
481            let prefix_end = format!("{}~", prefix); // Use tilde as it's after most printable chars
482            let filter = format!(
483                "(PartitionKey eq '{}') and (RowKey ge '{}') and (RowKey lt '{}')",
484                partition_key, prefix, prefix_end
485            );
486
487            // Note: We'll do TTL filtering client-side to avoid OData syntax issues
488            let filter_with_ttl = filter;
489
490            let query_options = EntityQueryOptions {
491                filter: Some(filter_with_ttl),
492                select: None,
493                top: Some((limit - total_fetched) as u32),
494            };
495
496            let response = self
497                .client
498                .query_entities(
499                    &self.resource_group_name,
500                    &self.account_name,
501                    &self.table_name,
502                    Some(query_options),
503                )
504                .await
505                .map_err(|e| {
506                    crate::error::map_cloud_client_error(
507                        e,
508                        format!("Failed to query entities with prefix '{}'", prefix),
509                        Some(prefix.to_string()),
510                    )
511                })?;
512
513            // Process entities from this partition
514            for entity in response.entities {
515                if total_fetched >= limit {
516                    break;
517                }
518
519                // Additional client-side TTL check for precision
520                if is_entity_expired(&entity) {
521                    continue; // Skip expired
522                }
523
524                let key = self.combine_key(&entity.partition_key, &entity.row_key);
525                let value = extract_value_from_entity(&entity)?;
526
527                all_items.push((key, value));
528                total_fetched += 1;
529            }
530
531            // If we hit the limit or have more data in this partition, encode cursor and return
532            if total_fetched >= limit || response.next_link.is_some() {
533                let next_cursor = self.encode_cursor(&CursorState {
534                    current_partition: partition_id,
535                    partition_continuation_token: response.next_link,
536                });
537                return Ok(ScanResult {
538                    items: all_items,
539                    next_cursor: Some(next_cursor),
540                });
541            }
542        }
543
544        // Scanned all partitions without hitting limit
545        Ok(ScanResult {
546            items: all_items,
547            next_cursor: None,
548        })
549    }
550}
alien_bindings/providers/kv/azure_table_storage.rs

alien_bindings/providers/kv/
azure_table_storage.rs
