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

alien_bindings/providers/kv/
azure_table_storage.rs
