pdk_data_storage_lib/
lib.rs

1// Copyright (c) 2026, Salesforce, Inc.,
2// All rights reserved.
3// For full license text, see the LICENSE.txt file
4
5//! Data Storage Library
6//!
7//! This library provides data storage functionality with support for:
8//!
9//! - Local and distributed data storage
10//! - Support for CAS (Compare-And-Swap) operations
11//! - Configurable storage modes (Always, Absent, CAS)
12//! - Asynchronous API for high-performance applications
13//!
14//! ## Features
15//! - `ll`: USE AT OWN RISK: low-level items that may change without notice. Exposes the underlying implementation of each storage type.
16//! - `experimental_datastorage_formats`: USE AT OWN RISK: experimental support for user-defined serialization formats.
17
18mod local;
19
20mod distributed;
21
22mod format;
23
24#[cfg(feature = "experimental_datastorage_formats")]
25pub use format::{Format, MessagePack};
26use std::marker::PhantomData;
27
28#[cfg(feature = "ll")]
29/// "Low Level" implementations of each storage type.
30pub mod ll {
31    /// In memory storage implementation.
32    pub mod local {
33        pub use crate::local::*;
34    }
35
36    /// External storage implementation.
37    pub mod distributed {
38        pub use crate::distributed::*;
39    }
40}
41
42use pdk_core::classy::extract::context::ConfigureContext;
43use pdk_core::classy::extract::{Extract, FromContext};
44use pdk_core::logger;
45use serde::{de::DeserializeOwned, Serialize};
46use std::rc::Rc;
47use thiserror::Error;
48use url::form_urlencoded;
49
50use crate::distributed::DistributedStorage;
51use crate::format::{DefaultFormat, FormatWrapper};
52use crate::local::LocalStorage;
53
54/// Defines the behavior for store operations.
55///
56/// This enum specifies how a store operation should behave when a key already exists
57/// in the storage system.
58#[derive(PartialEq, Eq, Debug, Clone)]
59pub enum StoreMode {
60    /// Indicates that the store operation should always succeed, overwriting any existing value.
61    ///
62    /// This mode is useful when you want to unconditionally update a value regardless
63    /// of its current state.
64    Always,
65    /// Indicates that the store operation should succeed only if no value was previously stored.
66    ///
67    /// This mode is useful for implementing "set if not exists" semantics, ensuring
68    /// that values are only written once.
69    Absent,
70    /// Indicates that the store operation should succeed only if the stored value matches the provided CAS.
71    ///
72    /// This mode is useful for implementing optimistic concurrency control, ensuring
73    /// that updates only succeed if the value hasn't been modified by another operation.
74    Cas(String),
75}
76
77/// Errors that can occur during data storage operations.
78///
79/// This enum represents all possible error conditions that can arise when
80/// performing data storage operations.
81#[derive(Debug, Error)]
82#[non_exhaustive]
83pub enum DataStorageError {
84    /// Indicates provided CAS value doesn't match current version of the stored value.
85    #[error("CAS mismatch.")]
86    CasMismatch,
87    #[error("Serialization error: {0}.")]
88    #[cfg(feature = "experimental_datastorage_formats")]
89    /// Indicates the requested operation failed due to a serialization error.
90    Format(Box<dyn std::error::Error>),
91    #[error("Serialization error: {0}.")]
92    /// Indicates the requested operation failed due to a serialization error.
93    Serialization(#[from] serde_fixint::Error),
94    /// Indicates the requested operation failed due to a CAS parse error.
95    #[error("CAS parse error: {0}.")]
96    CasParseError(#[from] std::num::ParseIntError),
97    /// Indicates the operation timed out.
98    #[error("Timeout.")]
99    Timeout,
100    /// Indicates an error occurred while performing an http client call.
101    #[error("HTTP Client Error.")]
102    HttpClient,
103    /// Indicates an unexpected error occurred.
104    #[error("Unexpected error: {0}.")]
105    Unexpected(String),
106}
107
108/// Errors that can occur when building data storage instances.
109///
110/// This enum represents error conditions that can arise when creating
111/// data storage instances using the builder pattern.
112#[derive(Debug, Error)]
113#[non_exhaustive]
114pub enum DataStorageBuilderError {
115    /// Local storage is required but not available in the current context.
116    #[error("Local storage not available")]
117    LocalStorageRequired,
118
119    /// Policy metadata is required but not available in the current context.
120    #[error("Policy metadata not available")]
121    MetadataRequired,
122}
123
124impl From<crate::local::LocalStorageError> for DataStorageError {
125    fn from(error: crate::local::LocalStorageError) -> Self {
126        match error {
127            crate::local::LocalStorageError::CasMismatch => DataStorageError::CasMismatch,
128            _ => DataStorageError::Unexpected(error.to_string()),
129        }
130    }
131}
132
133impl From<crate::distributed::DistributedStorageError> for DataStorageError {
134    fn from(error: crate::distributed::DistributedStorageError) -> Self {
135        match error {
136            crate::distributed::DistributedStorageError::CasMismatch => {
137                DataStorageError::CasMismatch
138            }
139            crate::distributed::DistributedStorageError::Timeout => DataStorageError::Timeout,
140            crate::distributed::DistributedStorageError::HttpClient(_) => {
141                DataStorageError::HttpClient
142            }
143            error => DataStorageError::Unexpected(error.to_string()),
144        }
145    }
146}
147
148/// A trait for data storage operations that can be implemented by different storage backends.
149///
150/// This trait defines the core interface for data storage operations, allowing
151/// implementations to use different storage backends (local, distributed) while
152/// providing a consistent API.
153#[allow(async_fn_in_trait)]
154pub trait DataStorage {
155    /// Returns all keys currently stored in the store.
156    async fn get_keys(&self) -> Result<Vec<String>, DataStorageError>;
157
158    /// Stores a serializable item for a given key using the provided [`StoreMode`].
159    async fn store<T: Serialize>(
160        &self,
161        key: &str,
162        mode: &StoreMode,
163        item: &T,
164    ) -> Result<(), DataStorageError>;
165
166    /// Retrieves and deserializes the value for a given key, returning the value and its CAS
167    /// string if present. Returns `Ok(None)` when the key does not exist.
168    async fn get<T: DeserializeOwned>(
169        &self,
170        key: &str,
171    ) -> Result<Option<(T, String)>, DataStorageError>;
172
173    /// Removes the item identified by the provided key from this storage instance.
174    async fn delete(&self, key: &str) -> Result<(), DataStorageError>;
175
176    /// Removes all items from this storage instance.
177    async fn delete_all(&self) -> Result<(), DataStorageError>;
178}
179
180/// A local data storage implementation that stores data in memory.
181///
182/// This implementation uses local shared data storage for high-performance
183/// in-memory operations within a single node.
184pub struct LocalDataStorage<F: format::Format = format::DefaultFormat> {
185    storage: crate::local::SharedData,
186    namespace: String,
187    format: FormatWrapper<F>,
188}
189
190impl<F: format::Format> LocalDataStorage<F> {
191    /// Creates a new local data storage instance using the give storage and namespace as unique identifier.
192    pub(crate) fn new(
193        storage: crate::local::SharedData,
194        namespace: String,
195        format: FormatWrapper<F>,
196    ) -> Self {
197        Self {
198            storage,
199            namespace,
200            format,
201        }
202    }
203
204    fn convert_store_mode(
205        &self,
206        mode: &StoreMode,
207    ) -> Result<crate::local::StoreMode, DataStorageError> {
208        match mode {
209            StoreMode::Always => Ok(crate::local::StoreMode::Always),
210            StoreMode::Absent => Ok(crate::local::StoreMode::Absent),
211            StoreMode::Cas(cas_str) => {
212                let cas: u32 = cas_str.parse()?;
213                Ok(crate::local::StoreMode::Cas(cas))
214            }
215        }
216    }
217
218    fn namespaced_key(&self, key: &str) -> String {
219        format!("{}:{}", self.namespace, key)
220    }
221
222    fn get_keys(&self) -> Result<Vec<String>, DataStorageError> {
223        let all_keys = self.storage.keys();
224        let namespace_prefix = format!("{}:", self.namespace);
225        Ok(all_keys
226            .into_iter()
227            .filter(|key| key.starts_with(&namespace_prefix))
228            .map(|key| {
229                key.strip_prefix(&namespace_prefix)
230                    .unwrap_or(&key)
231                    .to_string()
232            })
233            .collect())
234    }
235
236    fn store<T: Serialize>(
237        &self,
238        key: &str,
239        mode: &StoreMode,
240        item: &T,
241    ) -> Result<(), DataStorageError> {
242        let serialized = self.format.serialize(item)?;
243        let local_mode = self.convert_store_mode(mode)?;
244        let namespaced_key = self.namespaced_key(key);
245        self.storage.set(&namespaced_key, &serialized, local_mode)?;
246        Ok(())
247    }
248
249    fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<(T, String)>, DataStorageError> {
250        let namespaced_key = self.namespaced_key(key);
251        match self.storage.get(&namespaced_key)? {
252            Some((data, cas)) => {
253                let deserialized: T = self.format.deserialize(&data)?;
254                Ok(Some((deserialized, cas.to_string())))
255            }
256            None => Ok(None),
257        }
258    }
259
260    fn delete(&self, key: &str) -> Result<(), DataStorageError> {
261        let namespaced_key = self.namespaced_key(key);
262        self.storage.delete(&namespaced_key)?;
263        Ok(())
264    }
265
266    fn delete_all(&self) -> Result<(), DataStorageError> {
267        let all_keys = self.storage.keys();
268        let namespace_prefix = format!("{}:", self.namespace);
269        for key in all_keys {
270            if key.starts_with(&namespace_prefix) {
271                self.storage.delete(&key)?;
272            }
273        }
274        Ok(())
275    }
276
277    /// Returns a handle for synchronous (blocking) operations on this storage instance.
278    #[cfg(feature = "experimental_storage_sync")]
279    pub fn blocking(&self) -> &impl BlockingDataStorage {
280        self
281    }
282}
283
284impl<F: format::Format> DataStorage for LocalDataStorage<F> {
285    async fn get_keys(&self) -> Result<Vec<String>, DataStorageError> {
286        self.get_keys()
287    }
288
289    async fn store<T: Serialize>(
290        &self,
291        key: &str,
292        mode: &StoreMode,
293        item: &T,
294    ) -> Result<(), DataStorageError> {
295        self.store(key, mode, item)
296    }
297
298    async fn get<T: DeserializeOwned>(
299        &self,
300        key: &str,
301    ) -> Result<Option<(T, String)>, DataStorageError> {
302        self.get(key)
303    }
304
305    async fn delete(&self, key: &str) -> Result<(), DataStorageError> {
306        self.delete(key)
307    }
308
309    async fn delete_all(&self) -> Result<(), DataStorageError> {
310        self.delete_all()
311    }
312}
313
314/// Synchronous (blocking) storage operations.
315/// Obtained via [`LocalDataStorage::blocking`]. Use when `.await` is not available,
316#[cfg(feature = "experimental_storage_sync")]
317pub trait BlockingDataStorage {
318    /// Returns all keys currently stored in this namespace.
319    fn get_keys(&self) -> Result<Vec<String>, DataStorageError>;
320
321    /// Stores a serializable item for a given key using the provided [`StoreMode`].
322    fn store<T: Serialize>(
323        &self,
324        key: &str,
325        mode: &StoreMode,
326        item: &T,
327    ) -> Result<(), DataStorageError>;
328
329    /// Retrieves and deserializes the value for a given key, returning the value and its CAS
330    /// string if present. Returns `Ok(None)` when the key does not exist.
331    fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<(T, String)>, DataStorageError>;
332
333    /// Removes the item identified by the provided key from this storage instance.
334    fn delete(&self, key: &str) -> Result<(), DataStorageError>;
335
336    /// Removes all items from this storage instance.
337    fn delete_all(&self) -> Result<(), DataStorageError>;
338}
339
340#[cfg(feature = "experimental_storage_sync")]
341impl<F: format::Format> BlockingDataStorage for LocalDataStorage<F> {
342    fn get_keys(&self) -> Result<Vec<String>, DataStorageError> {
343        self.get_keys()
344    }
345
346    fn store<T: Serialize>(
347        &self,
348        key: &str,
349        mode: &StoreMode,
350        item: &T,
351    ) -> Result<(), DataStorageError> {
352        self.store(key, mode, item)
353    }
354
355    fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<(T, String)>, DataStorageError> {
356        self.get(key)
357    }
358
359    fn delete(&self, key: &str) -> Result<(), DataStorageError> {
360        self.delete(key)
361    }
362
363    fn delete_all(&self) -> Result<(), DataStorageError> {
364        self.delete_all()
365    }
366}
367
368/// A distributed data storage implementation that stores data across multiple nodes.
369pub struct RemoteDataStorage<F: format::Format = format::DefaultFormat> {
370    storage: Rc<crate::distributed::DistributedStorageClient>,
371    sanitized_store: String,
372    sanitized_partition: String,
373    ttl_millis: u32,
374    format: FormatWrapper<F>,
375}
376
377impl<F: format::Format> RemoteDataStorage<F> {
378    /// Creates a high level data storage instance for the given `store`/`partition`,
379    /// using `storage` and a default TTL in milliseconds.
380    pub(crate) fn new(
381        storage: Rc<crate::distributed::DistributedStorageClient>,
382        store: String,
383        partition: String,
384        ttl_millis: u32,
385        format: FormatWrapper<F>,
386    ) -> Self {
387        // Sanitize store and partition names using form_urlencoded
388        let sanitized_store = form_urlencoded::byte_serialize(store.as_bytes()).collect();
389        let sanitized_partition = form_urlencoded::byte_serialize(partition.as_bytes()).collect();
390        Self {
391            storage,
392            sanitized_store,
393            sanitized_partition,
394            ttl_millis,
395            format,
396        }
397    }
398
399    fn convert_store_mode(&self, mode: &StoreMode) -> crate::distributed::StoreMode {
400        match mode {
401            StoreMode::Always => crate::distributed::StoreMode::Always,
402            StoreMode::Absent => crate::distributed::StoreMode::Absent,
403            StoreMode::Cas(cas_str) => crate::distributed::StoreMode::Cas(cas_str.clone()),
404        }
405    }
406
407    fn sanitize_key(&self, key: &str) -> String {
408        form_urlencoded::byte_serialize(key.as_bytes()).collect()
409    }
410}
411
412impl<F: format::Format> DataStorage for RemoteDataStorage<F> {
413    async fn get_keys(&self) -> Result<Vec<String>, DataStorageError> {
414        // Try to get keys, if store doesn't exist, return empty list
415        match self
416            .storage
417            .get_keys(&self.sanitized_store, &self.sanitized_partition)
418            .await
419        {
420            Ok(keys) => {
421                // Decode the keys to return them in the same format as they were stored
422                let decoded_keys: Vec<String> = keys
423                    .into_iter()
424                    .filter_map(|encoded_key| {
425                        let decoded = form_urlencoded::parse(encoded_key.as_bytes())
426                            .next()
427                            .map(|(key, _)| key.into_owned());
428
429                        if decoded.is_none() {
430                            logger::debug!("Key not URL-encoded or decode failed: {encoded_key}");
431                        }
432
433                        decoded
434                    })
435                    .collect();
436
437                Ok(decoded_keys)
438            }
439            Err(e) => {
440                logger::warn!("Error getting keys: {e}");
441                Ok(vec![])
442            }
443        }
444    }
445
446    async fn store<T: Serialize>(
447        &self,
448        key: &str,
449        mode: &StoreMode,
450        item: &T,
451    ) -> Result<(), DataStorageError> {
452        let serialized = self.format.serialize(item)?;
453        let distributed_mode = self.convert_store_mode(mode);
454        let sanitized_key = self.sanitize_key(key);
455
456        // Try to store first
457        match self
458            .storage
459            .store(
460                &self.sanitized_store,
461                &self.sanitized_partition,
462                &sanitized_key,
463                &distributed_mode,
464                &serialized,
465            )
466            .await
467        {
468            Ok(()) => Ok(()),
469            Err(crate::distributed::DistributedStorageError::StoreNotFound) => {
470                // Store doesn't exist, create it and retry
471                let store = crate::distributed::Store::new(
472                    self.sanitized_store.clone(),
473                    Some(self.ttl_millis),
474                    None,
475                );
476
477                // Try to create the store, ignore if it already exists
478                if let Err(e) = self.storage.upsert_store(&store).await {
479                    logger::warn!("Error creating store: {e}");
480                }
481
482                // Try storing again
483                self.storage
484                    .store(
485                        &self.sanitized_store,
486                        &self.sanitized_partition,
487                        &sanitized_key,
488                        &distributed_mode,
489                        &serialized,
490                    )
491                    .await?;
492                Ok(())
493            }
494            Err(e) => Err(e.into()), // Other errors, propagate them
495        }
496    }
497
498    async fn get<T: DeserializeOwned>(
499        &self,
500        key: &str,
501    ) -> Result<Option<(T, String)>, DataStorageError> {
502        let sanitized_key = self.sanitize_key(key);
503        match self
504            .storage
505            .get(
506                &self.sanitized_store,
507                &self.sanitized_partition,
508                &sanitized_key,
509            )
510            .await
511        {
512            Ok((data, cas)) => {
513                let deserialized: T = self.format.deserialize(&data)?;
514                Ok(Some((deserialized, cas)))
515            }
516            Err(crate::distributed::DistributedStorageError::StoreNotFound) => {
517                logger::debug!("Store not found for key {key}, returning None");
518                Ok(None)
519            }
520            Err(crate::distributed::DistributedStorageError::KeyNotFound) => {
521                logger::debug!("Key not found: {key}");
522                Ok(None)
523            }
524            Err(e) => {
525                logger::error!("Error getting value for key {key}: {e:?}");
526                Err(e.into())
527            }
528        }
529    }
530
531    async fn delete(&self, key: &str) -> Result<(), DataStorageError> {
532        // Try to delete, ignore if store doesn't exist
533        let sanitized_key = self.sanitize_key(key);
534        if let Err(e) = self
535            .storage
536            .delete(
537                &self.sanitized_store,
538                &self.sanitized_partition,
539                &sanitized_key,
540            )
541            .await
542        {
543            logger::warn!("Error deleting key {key}: {e}");
544        }
545        Ok(())
546    }
547
548    async fn delete_all(&self) -> Result<(), DataStorageError> {
549        // Try to delete partition, ignore if store doesn't exist
550        if let Err(e) = self
551            .storage
552            .delete_partition(&self.sanitized_store, &self.sanitized_partition)
553            .await
554        {
555            logger::warn!("Error deleting partition: {e}");
556        }
557        Ok(())
558    }
559}
560
561/// Builder for creating data storage instances.
562///
563/// Provides methods to create local and distributed storage instances with
564/// configurable settings. The builder pattern allows for flexible configuration
565/// of storage behavior.
566///
567/// # Examples
568///
569/// ```rust
570/// # use data_storage_lib::{DataStorageBuilder, DataStorage};
571/// # async fn example(builder: DataStorageBuilder) {
572/// // Create a local storage instance
573/// let local_storage = builder.local("my-local-storage");
574///
575/// // Create a distributed storage instance with 60-second TTL
576/// let remote_storage = builder.remote("my-remote-storage", 60000);
577/// # }
578/// ```
579pub struct DataStorageBuilder<F: format::Format = format::DefaultFormat> {
580    prefix: String,
581    shared_data: Rc<crate::local::SharedData>,
582    distributed_storage: Option<Rc<crate::distributed::DistributedStorageClient>>,
583    format_wrapper: FormatWrapper<F>,
584}
585
586impl<F: format::Format + Clone> Clone for DataStorageBuilder<F> {
587    fn clone(&self) -> Self {
588        Self {
589            prefix: self.prefix.clone(),
590            shared_data: Rc::clone(&self.shared_data),
591            distributed_storage: self.distributed_storage.clone(),
592            format_wrapper: self.format_wrapper.clone(),
593        }
594    }
595}
596
597/// DataStorageBuilder can be injected in your configuration function.
598/// ```rust
599/// #[entrypoint]
600/// async fn configure(
601///     launcher: Launcher,
602///     store_builder: DataStorageBuilder,
603///     Configuration(configuration): Configuration,
604/// ) -> anyhow::Result<()> {
605/// }
606/// ```
607impl FromContext<ConfigureContext> for DataStorageBuilder<DefaultFormat> {
608    type Error = DataStorageBuilderError;
609
610    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
611        // Extract local storage (required)
612        let shared_data: crate::local::SharedData = context
613            .extract()
614            .map_err(|_| DataStorageBuilderError::LocalStorageRequired)?;
615        // Extract distributed storage (optional - will be None if not available)
616        let distributed_storage: Result<crate::distributed::DistributedStorageClient, _> =
617            context.extract();
618        // Extract metadata for policy isolation
619        let metadata: pdk_core::policy_context::api::Metadata = context
620            .extract()
621            .map_err(|_| DataStorageBuilderError::MetadataRequired)?;
622
623        let prefix = format!(
624            "isolated-storage-{}-{}",
625            metadata.policy_metadata.policy_name, metadata.policy_metadata.policy_namespace
626        );
627
628        pdk_core::logger::debug!(
629            "DataStorageBuilder: creating prefix '{}' for policy '{}' in namespace '{}'",
630            prefix,
631            metadata.policy_metadata.policy_name,
632            metadata.policy_metadata.policy_namespace
633        );
634
635        Ok(DataStorageBuilder {
636            prefix,
637            shared_data: Rc::new(shared_data),
638            distributed_storage: distributed_storage.ok().map(Rc::new),
639            format_wrapper: FormatWrapper::Default(DefaultFormat::default(), PhantomData),
640        })
641    }
642}
643
644impl<F: format::Format + Clone> DataStorageBuilder<F> {
645    /// Indicates that storage should not be isolated to a single policy.
646    /// The resulting state will be shared across policy instances that use
647    /// the same storage ID.
648    pub fn shared(mut self) -> Self {
649        self.prefix = "shared-storage".to_string();
650        self
651    }
652
653    /// Sets the serialization format.
654    #[cfg(feature = "experimental_datastorage_formats")]
655    pub fn format<H: Format + Clone>(self, format: H) -> DataStorageBuilder<H> {
656        DataStorageBuilder {
657            prefix: self.prefix,
658            shared_data: self.shared_data,
659            distributed_storage: self.distributed_storage,
660            format_wrapper: FormatWrapper::Custom(format),
661        }
662    }
663
664    /// Creates a local data storage instance identified by the key.
665    pub fn local<T: Into<String>>(&self, key: T) -> LocalDataStorage<F> {
666        let key_str = key.into();
667        // Create a truly unique namespace by combining prefix and key
668        // This ensures each policy gets its own isolated storage
669        let namespace = format!("{}-{}", self.prefix, key_str);
670
671        pdk_core::logger::debug!(
672            "DataStorageBuilder::local: creating namespace '{}' with prefix '{}' and key '{}'",
673            namespace,
674            self.prefix,
675            key_str
676        );
677
678        LocalDataStorage::new(
679            (*self.shared_data).clone(),
680            namespace,
681            self.format_wrapper.clone(),
682        )
683    }
684
685    /// Creates a distributed data storage instance identified by the provided key and ttl in milliseconds.
686    ///
687    /// **Note**: To make use of the remote mode of the data storage your Flex Gateway must have [Shared Storage](https://docs.mulesoft.com/gateway/latest/flex-conn-shared-storage-config) configured.
688    ///
689    /// # Panics
690    ///
691    /// Panics if distributed storage is not available in the current context.
692    /// Ensure that distributed storage is properly configured before calling this method.
693    pub fn remote<T: Into<String>>(&self, key: T, ttl_millis: u32) -> RemoteDataStorage<F> {
694        let key_str = key.into();
695        let storage = self
696            .distributed_storage
697            .as_ref()
698            .expect("Distributed storage not available - check if it's configured");
699        RemoteDataStorage::new(
700            Rc::clone(storage),
701            key_str.clone(),
702            key_str,
703            ttl_millis,
704            self.format_wrapper.clone(),
705        )
706    }
707}
pdk_data_storage_lib/lib.rs

pdk_data_storage_lib/
lib.rs
