alien_bindings/providers/kv/
azure_table_storage.rs

1use crate::error::{ErrorData, Result};
2use crate::traits::{Binding, Kv, KvEntry, PutCondition, 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::{decode_version, encode_version, 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
88fn entity_expiration_millis(entity: &TableEntity) -> Option<i64> {
89    entity
90        .properties
91        .get("ExpiresAt")
92        .and_then(Value::as_str)
93        .and_then(|value| DateTime::parse_from_rfc3339(value).ok())
94        .map(|value| value.timestamp_millis())
95}
96
97fn entity_etag(entity: &TableEntity) -> Option<String> {
98    entity
99        .properties
100        .get("odata.etag")
101        .and_then(Value::as_str)
102        .map(ToString::to_string)
103}
104
105/// Cursor state for pagination across partitions
106#[derive(Serialize, Deserialize)]
107#[serde(rename_all = "camelCase")]
108struct CursorState {
109    version: u8,
110    prefix: String,
111    current_partition: u32,
112    continuation: Option<EntityQueryContinuation>,
113}
114
115/// Azure Table Storage implementation of the KV trait
116pub struct AzureTableStorageKv {
117    client: AzureTableStorageClient,
118    resource_group_name: String,
119    account_name: String,
120    table_name: String,
121    num_partitions: u32,
122}
123
124impl Debug for AzureTableStorageKv {
125    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
126        f.debug_struct("AzureTableStorageKv")
127            .field("resource_group_name", &self.resource_group_name)
128            .field("account_name", &self.account_name)
129            .field("table_name", &self.table_name)
130            .field("num_partitions", &self.num_partitions)
131            .finish()
132    }
133}
134
135impl AzureTableStorageKv {
136    pub fn new(
137        client: AzureTableStorageClient,
138        resource_group_name: String,
139        account_name: String,
140        table_name: String,
141    ) -> Self {
142        Self {
143            client,
144            resource_group_name,
145            account_name,
146            table_name,
147            num_partitions: 16, // 16 partitions for load distribution
148        }
149    }
150
151    /// Creates a hash bucket for load distribution
152    fn hash_bucket(&self, key: &str) -> u32 {
153        use std::collections::hash_map::DefaultHasher;
154        use std::hash::{Hash, Hasher};
155
156        let mut hasher = DefaultHasher::new();
157        key.hash(&mut hasher);
158        hasher.finish() as u32 % self.num_partitions
159    }
160
161    /// Splits key into partition key and row key
162    /// Attempt to take over a logically-expired row after a conditional
163    /// insert conflict. Returns `Ok(true)` only when THIS caller replaced
164    /// the expired row (If-Match on the read ETag); a live row, a lost
165    /// race, or a row deleted in between all resolve to `Ok(false)`.
166    async fn try_take_over_expired(
167        &self,
168        key: &str,
169        entity: &alien_azure_clients::azure::tables::TableEntity,
170    ) -> Result<bool> {
171        use alien_client_core::ErrorData as CloudErrorData;
172
173        let (partition_key, row_key) = self.split_key(key);
174        let existing = match self
175            .client
176            .get_entity(
177                &self.resource_group_name,
178                &self.account_name,
179                &self.table_name,
180                &partition_key,
181                &row_key,
182                None,
183            )
184            .await
185        {
186            Ok(existing) => existing,
187            // Deleted between the conflict and this read: treat as lost —
188            // the caller's next attempt sees a clean slate.
189            Err(e)
190                if matches!(
191                    e.error.as_ref(),
192                    Some(CloudErrorData::RemoteResourceNotFound { .. })
193                ) =>
194            {
195                return Ok(false);
196            }
197            Err(e) => {
198                return Err(crate::error::map_cloud_client_error(
199                    e,
200                    format!("Failed to read existing entity for key '{}'", key),
201                    Some(key.to_string()),
202                ));
203            }
204        };
205
206        if !is_entity_expired(&existing) {
207            return Ok(false);
208        }
209        let Some(etag) = existing
210            .properties
211            .get("odata.etag")
212            .and_then(|value| value.as_str())
213        else {
214            // No ETag on the read entity — cannot replace safely.
215            return Ok(false);
216        };
217
218        match self
219            .client
220            .update_entity(
221                &self.resource_group_name,
222                &self.account_name,
223                &self.table_name,
224                &partition_key,
225                &row_key,
226                entity,
227                Some(alien_azure_clients::azure::tables::ETag::from(etag)),
228            )
229            .await
230        {
231            Ok(_) => Ok(true),
232            // 409/412: someone else replaced or deleted it first.
233            Err(e)
234                if matches!(
235                    e.error.as_ref(),
236                    Some(CloudErrorData::RemoteResourceConflict { .. })
237                        | Some(CloudErrorData::RemoteResourceNotFound { .. })
238                ) =>
239            {
240                Ok(false)
241            }
242            Err(e) => Err(crate::error::map_cloud_client_error(
243                e,
244                format!("Failed to take over expired entity for key '{}'", key),
245                Some(key.to_string()),
246            )),
247        }
248    }
249
250    fn split_key(&self, key: &str) -> (String, String) {
251        // Use hash-based partitioning for load distribution
252        let partition_key = format!("p{}", self.hash_bucket(key));
253        (partition_key, key.to_string())
254    }
255
256    /// Combines partition key and row key back to original key
257    fn combine_key(&self, _partition_key: &str, row_key: &str) -> String {
258        row_key.to_string() // Row key contains the original key
259    }
260
261    /// Encodes cursor state as base64url JSON for safe HTTP transmission
262    fn encode_cursor(&self, state: &CursorState) -> Result<String> {
263        let json =
264            serde_json::to_vec(state)
265                .into_alien_error()
266                .context(ErrorData::InvalidInput {
267                    operation_context: "Azure Table Storage KV cursor encoding".to_string(),
268                    details: "Failed to serialize cursor state".to_string(),
269                    field_name: Some("cursor".to_string()),
270                })?;
271        Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(json))
272    }
273
274    /// Decodes cursor state from base64url JSON
275    fn decode_cursor(&self, prefix: &str, cursor: &str) -> Result<CursorState> {
276        let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
277            .decode(cursor)
278            .into_alien_error()
279            .context(ErrorData::InvalidInput {
280                operation_context: "Azure Table Storage KV cursor decoding".to_string(),
281                details: "Invalid cursor encoding".to_string(),
282                field_name: Some("cursor".to_string()),
283            })?;
284        let json =
285            String::from_utf8(decoded)
286                .into_alien_error()
287                .context(ErrorData::InvalidInput {
288                    operation_context: "Azure Table Storage KV cursor decoding".to_string(),
289                    details: "Invalid cursor UTF-8".to_string(),
290                    field_name: Some("cursor".to_string()),
291                })?;
292        let state: CursorState =
293            serde_json::from_str(&json)
294                .into_alien_error()
295                .context(ErrorData::InvalidInput {
296                    operation_context: "Azure Table Storage KV cursor decoding".to_string(),
297                    details: "Invalid cursor JSON".to_string(),
298                    field_name: Some("cursor".to_string()),
299                })?;
300        if state.version != 1
301            || state.prefix != prefix
302            || state.current_partition >= self.num_partitions
303        {
304            return Err(AlienError::new(ErrorData::InvalidInput {
305                operation_context: "Azure Table Storage KV cursor validation".to_string(),
306                details: "Cursor does not belong to this prefix scan".to_string(),
307                field_name: Some("cursor".to_string()),
308            }));
309        }
310        Ok(state)
311    }
312}
313
314impl Binding for AzureTableStorageKv {}
315
316#[async_trait]
317impl Kv for AzureTableStorageKv {
318    async fn get(&self, key: &str) -> Result<Option<KvEntry>> {
319        validate_key(key)?;
320
321        let (partition_key, row_key) = self.split_key(key);
322
323        match self
324            .client
325            .get_entity(
326                &self.resource_group_name,
327                &self.account_name,
328                &self.table_name,
329                &partition_key,
330                &row_key,
331                None,
332            )
333            .await
334        {
335            Ok(entity) => {
336                // Check if TTL has expired (client-side filtering)
337                if is_entity_expired(&entity) {
338                    return Ok(None); // Expired
339                }
340
341                let value = extract_value_from_entity(&entity)?;
342                let etag = entity_etag(&entity).ok_or_else(|| {
343                    AlienError::new(ErrorData::UnexpectedResponseFormat {
344                        provider: "azure".to_string(),
345                        binding_name: "table-storage".to_string(),
346                        field: "odata.etag".to_string(),
347                        response_json: serde_json::to_string(&entity).unwrap_or_default(),
348                    })
349                })?;
350                Ok(Some(KvEntry {
351                    key: key.to_string(),
352                    value,
353                    version: encode_version(key, etag, entity_expiration_millis(&entity))?,
354                }))
355            }
356            Err(e) => {
357                use alien_client_core::ErrorData as CloudErrorData;
358                match e.error.as_ref() {
359                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(None),
360                    _ => Err(crate::error::map_cloud_client_error(
361                        e,
362                        format!("Failed to get entity for key '{}'", key),
363                        Some(key.to_string()),
364                    )),
365                }
366            }
367        }
368    }
369
370    async fn put(&self, key: &str, value: Vec<u8>, options: Option<PutOptions>) -> Result<bool> {
371        validate_key(key)?;
372        validate_value(&value)?;
373
374        let options = options.unwrap_or_default();
375        let (partition_key, row_key) = self.split_key(key);
376
377        let expires_at = options.ttl.map(|d| Utc::now() + d);
378        let entity =
379            create_table_entity(partition_key.clone(), row_key.clone(), &value, expires_at);
380
381        if matches!(options.condition, PutCondition::Absent) {
382            match self
383                .client
384                .insert_entity(
385                    &self.resource_group_name,
386                    &self.account_name,
387                    &self.table_name,
388                    &entity,
389                )
390                .await
391            {
392                Ok(_) => Ok(true),
393                Err(e) => {
394                    use alien_client_core::ErrorData as CloudErrorData;
395                    match e.error.as_ref() {
396                        Some(CloudErrorData::RemoteResourceConflict { .. }) => {
397                            // Expired rows count as ABSENT, matching the
398                            // local provider's atomic takeover: Table
399                            // Storage has no server-side TTL, so without
400                            // this an expired row (e.g. a dead command
401                            // lease) blocks conditional puts forever. The
402                            // read + If-Match replace is race-safe: a
403                            // concurrent taker changes the ETag and this
404                            // update loses with a conflict.
405                            self.try_take_over_expired(key, &entity).await
406                        }
407                        _ => Err(crate::error::map_cloud_client_error(
408                            e,
409                            format!("Failed to insert entity for key '{}'", key),
410                            Some(key.to_string()),
411                        )),
412                    }
413                }
414            }
415        } else if let PutCondition::Version(ref version) = options.condition {
416            let expected = decode_version(key, version)?;
417            if expected.expired {
418                return Ok(false);
419            }
420            match self
421                .client
422                .update_entity(
423                    &self.resource_group_name,
424                    &self.account_name,
425                    &self.table_name,
426                    &partition_key,
427                    &row_key,
428                    &entity,
429                    Some(alien_azure_clients::azure::tables::ETag::from(
430                        expected.backend_version,
431                    )),
432                )
433                .await
434            {
435                Ok(_) => Ok(true),
436                Err(error)
437                    if matches!(
438                        error.error.as_ref(),
439                        Some(alien_client_core::ErrorData::RemoteResourceConflict { .. })
440                            | Some(alien_client_core::ErrorData::RemoteResourceNotFound { .. })
441                    ) =>
442                {
443                    Ok(false)
444                }
445                Err(error) => Err(crate::error::map_cloud_client_error(
446                    error,
447                    format!("Failed to conditionally update entity for key '{}'", key),
448                    Some(key.to_string()),
449                )),
450            }
451        } else {
452            // Insert Or Replace (upsert) - matches Azure REST API terminology
453            self.client
454                .insert_or_replace_entity(
455                    &self.resource_group_name,
456                    &self.account_name,
457                    &self.table_name,
458                    &partition_key,
459                    &row_key,
460                    &entity,
461                )
462                .await
463                .map_err(|e| {
464                    crate::error::map_cloud_client_error(
465                        e,
466                        format!("Failed to upsert entity for key '{}'", key),
467                        Some(key.to_string()),
468                    )
469                })?;
470            Ok(true)
471        }
472    }
473
474    async fn delete(&self, key: &str, if_version: Option<&str>) -> Result<bool> {
475        validate_key(key)?;
476
477        let (partition_key, row_key) = self.split_key(key);
478
479        let etag = if let Some(version) = if_version {
480            let expected = decode_version(key, version)?;
481            if expected.expired {
482                return Ok(false);
483            }
484            Some(alien_azure_clients::azure::tables::ETag::from(
485                expected.backend_version,
486            ))
487        } else {
488            None
489        };
490
491        match self
492            .client
493            .delete_entity(
494                &self.resource_group_name,
495                &self.account_name,
496                &self.table_name,
497                &partition_key,
498                &row_key,
499                etag,
500            )
501            .await
502        {
503            Ok(_) => Ok(true),
504            Err(e) => {
505                use alien_client_core::ErrorData as CloudErrorData;
506                match e.error.as_ref() {
507                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(if_version.is_none()),
508                    Some(CloudErrorData::RemoteResourceConflict { .. }) if if_version.is_some() => {
509                        Ok(false)
510                    }
511                    _ => Err(crate::error::map_cloud_client_error(
512                        e,
513                        format!("Failed to delete entity for key '{}'", key),
514                        Some(key.to_string()),
515                    )),
516                }
517            }
518        }
519    }
520
521    async fn exists(&self, key: &str) -> Result<bool> {
522        validate_key(key)?;
523
524        let (partition_key, row_key) = self.split_key(key);
525
526        match self
527            .client
528            .get_entity(
529                &self.resource_group_name,
530                &self.account_name,
531                &self.table_name,
532                &partition_key,
533                &row_key,
534                None,
535            )
536            .await
537        {
538            Ok(entity) => {
539                // Check TTL expiry
540                Ok(!is_entity_expired(&entity))
541            }
542            Err(e) => {
543                use alien_client_core::ErrorData as CloudErrorData;
544                match e.error.as_ref() {
545                    Some(CloudErrorData::RemoteResourceNotFound { .. }) => Ok(false),
546                    _ => Err(crate::error::map_cloud_client_error(
547                        e,
548                        format!("Failed to check existence of entity for key '{}'", key),
549                        Some(key.to_string()),
550                    )),
551                }
552            }
553        }
554    }
555
556    async fn scan_prefix(
557        &self,
558        prefix: &str,
559        limit: Option<usize>,
560        cursor: Option<String>,
561    ) -> Result<ScanResult> {
562        validate_key(prefix)?;
563        let limit = limit.unwrap_or(1000);
564        let initial = cursor
565            .as_deref()
566            .map(|cursor| self.decode_cursor(prefix, cursor))
567            .transpose()?
568            .unwrap_or(CursorState {
569                version: 1,
570                prefix: prefix.to_string(),
571                current_partition: 0,
572                continuation: None,
573            });
574        if limit == 0 {
575            return Ok(ScanResult {
576                items: Vec::new(),
577                next_cursor: cursor,
578            });
579        }
580
581        let mut items = Vec::with_capacity(limit);
582        let mut partition_id = initial.current_partition;
583        let mut continuation = initial.continuation;
584
585        while partition_id < self.num_partitions {
586            let partition_key = format!("p{}", partition_id);
587
588            // Build filter with BOTH PartitionKey and RowKey conditions
589            // Use a range query approach that's compatible with Azure Table Storage
590            let prefix_end = format!("{}~", prefix); // Use tilde as it's after most printable chars
591            let filter = format!(
592                "(PartitionKey eq '{}') and (RowKey ge '{}') and (RowKey lt '{}')",
593                partition_key, prefix, prefix_end
594            );
595
596            // Note: We'll do TTL filtering client-side to avoid OData syntax issues
597            let filter_with_ttl = filter;
598
599            let query_options = EntityQueryOptions {
600                filter: Some(filter_with_ttl),
601                select: None,
602                top: Some(u32::try_from(limit - items.len()).unwrap_or(u32::MAX)),
603                continuation: continuation.clone(),
604            };
605
606            let response = self
607                .client
608                .query_entities(
609                    &self.resource_group_name,
610                    &self.account_name,
611                    &self.table_name,
612                    Some(query_options),
613                )
614                .await
615                .map_err(|e| {
616                    crate::error::map_cloud_client_error(
617                        e,
618                        format!("Failed to query entities with prefix '{}'", prefix),
619                        Some(prefix.to_string()),
620                    )
621                })?;
622
623            for entity in response.entities {
624                if is_entity_expired(&entity) {
625                    continue;
626                }
627
628                let key = self.combine_key(&entity.partition_key, &entity.row_key);
629                let value = extract_value_from_entity(&entity)?;
630
631                if let Some(etag) = entity_etag(&entity) {
632                    items.push(KvEntry {
633                        version: encode_version(&key, etag, entity_expiration_millis(&entity))?,
634                        key,
635                        value,
636                    });
637                } else if let Some(entry) = self.get(&key).await? {
638                    items.push(entry);
639                }
640            }
641
642            if items.len() == limit {
643                let (next_partition, next_continuation) = match response.continuation {
644                    Some(continuation) => (partition_id, Some(continuation)),
645                    None if partition_id + 1 < self.num_partitions => (partition_id + 1, None),
646                    None => {
647                        return Ok(ScanResult {
648                            items,
649                            next_cursor: None,
650                        });
651                    }
652                };
653                return Ok(ScanResult {
654                    items,
655                    next_cursor: Some(self.encode_cursor(&CursorState {
656                        version: 1,
657                        prefix: prefix.to_string(),
658                        current_partition: next_partition,
659                        continuation: next_continuation,
660                    })?),
661                });
662            }
663
664            if let Some(next) = response.continuation {
665                continuation = Some(next);
666            } else {
667                partition_id += 1;
668                continuation = None;
669            }
670        }
671
672        Ok(ScanResult {
673            items,
674            next_cursor: None,
675        })
676    }
677}
alien_bindings/providers/kv/azure_table_storage.rs

alien_bindings/providers/kv/
azure_table_storage.rs
