pdk-data-storage-lib 1.6.0

// Copyright (c) 2025, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! Data Storage Library
//!
//! This library provides data storage functionality with support for:
//!
//! - Local and distributed data storage
//! - Support for CAS (Compare-And-Swap) operations
//! - Configurable storage modes (Always, Absent, CAS)
//! - Asynchronous API for high-performance applications
//!
//! ## Features
//! - `ll`: USE AT OWN RISK: low-level items that may change without notice. Exposes the underlying implementation of each storage type.

mod local;

mod distributed;

#[cfg(feature = "ll")]
/// "Low Level" implementations of each storage type.
pub mod ll {
    /// In memory storage implementation.
    pub mod local {
        pub use crate::local::*;
    }

    /// External storage implementation.
    pub mod distributed {
        pub use crate::distributed::*;
    }
}

use pdk_core::classy::extract::context::ConfigureContext;
use pdk_core::classy::extract::{Extract, FromContext};
use pdk_core::logger;
use serde::{de::DeserializeOwned, Serialize};
use std::rc::Rc;
use thiserror::Error;
use url::form_urlencoded;

use crate::distributed::DistributedStorage;
use crate::local::LocalStorage;

/// Defines the behavior for store operations.
///
/// This enum specifies how a store operation should behave when a key already exists
/// in the storage system.
#[derive(PartialEq, Eq, Debug, Clone)]
pub enum StoreMode {
    /// Indicates that the store operation should always succeed, overwriting any existing value.
    ///
    /// This mode is useful when you want to unconditionally update a value regardless
    /// of its current state.
    Always,
    /// Indicates that the store operation should succeed only if no value was previously stored.
    ///
    /// This mode is useful for implementing "set if not exists" semantics, ensuring
    /// that values are only written once.
    Absent,
    /// Indicates that the store operation should succeed only if the stored value matches the provided CAS.
    ///
    /// This mode is useful for implementing optimistic concurrency control, ensuring
    /// that updates only succeed if the value hasn't been modified by another operation.
    Cas(String),
}

/// Errors that can occur during data storage operations.
///
/// This enum represents all possible error conditions that can arise when
/// performing data storage operations.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum DataStorageError {
    /// Indicates provided CAS value doesn't match current version of the stored value.
    #[error("CAS mismatch.")]
    CasMismatch,
    /// Indicates the requested operation failed due to a serialization error.
    #[error("Serialization error: {0}.")]
    Serialization(#[from] bincode::Error),
    /// Indicates the requested operation failed due to a CAS parse error.
    #[error("CAS parse error: {0}.")]
    CasParseError(#[from] std::num::ParseIntError),
    /// Indicates the operation timed out.
    #[error("Timeout.")]
    Timeout,
    /// Indicates an error occurred while performing an http client call.
    #[error("HTTP Client Error.")]
    HttpClient,
    /// Indicates an unexpected error occurred.
    #[error("Unexpected error: {0}.")]
    Unexpected(String),
}

/// Errors that can occur when building data storage instances.
///
/// This enum represents error conditions that can arise when creating
/// data storage instances using the builder pattern.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum DataStorageBuilderError {
    /// Local storage is required but not available in the current context.
    #[error("Local storage not available")]
    LocalStorageRequired,

    /// Policy metadata is required but not available in the current context.
    #[error("Policy metadata not available")]
    MetadataRequired,
}

impl From<crate::local::LocalStorageError> for DataStorageError {
    fn from(error: crate::local::LocalStorageError) -> Self {
        match error {
            crate::local::LocalStorageError::CasMismatch => DataStorageError::CasMismatch,
            _ => DataStorageError::Unexpected(error.to_string()),
        }
    }
}

impl From<crate::distributed::DistributedStorageError> for DataStorageError {
    fn from(error: crate::distributed::DistributedStorageError) -> Self {
        match error {
            crate::distributed::DistributedStorageError::CasMismatch => {
                DataStorageError::CasMismatch
            }
            crate::distributed::DistributedStorageError::Timeout => DataStorageError::Timeout,
            crate::distributed::DistributedStorageError::HttpClient(_) => {
                DataStorageError::HttpClient
            }
            error => DataStorageError::Unexpected(error.to_string()),
        }
    }
}

/// A trait for data storage operations that can be implemented by different storage backends.
///
/// This trait defines the core interface for data storage operations, allowing
/// implementations to use different storage backends (local, distributed) while
/// providing a consistent API.
#[allow(async_fn_in_trait)]
pub trait DataStorage {
    /// Returns all keys currently stored in the store.
    async fn get_keys(&self) -> Result<Vec<String>, DataStorageError>;

    /// Stores a serializable item for a given key using the provided [`StoreMode`].
    async fn store<T: Serialize>(
        &self,
        key: &str,
        mode: &StoreMode,
        item: &T,
    ) -> Result<(), DataStorageError>;

    /// Retrieves and deserializes the value for a given key, returning the value and its CAS
    /// string if present. Returns `Ok(None)` when the key does not exist.
    async fn get<T: DeserializeOwned>(
        &self,
        key: &str,
    ) -> Result<Option<(T, String)>, DataStorageError>;

    /// Removes the item identified by the provided key from this storage instance.
    async fn delete(&self, key: &str) -> Result<(), DataStorageError>;

    /// Removes all items from this storage instance.
    async fn delete_all(&self) -> Result<(), DataStorageError>;
}

/// A local data storage implementation that stores data in memory.
///
/// This implementation uses local shared data storage for high-performance
/// in-memory operations within a single node.
pub struct LocalDataStorage {
    storage: crate::local::SharedData,
    namespace: String,
}

impl LocalDataStorage {
    /// Creates a new local data storage instance using the give storage and namespace as unique identifier.
    pub(crate) fn new(storage: crate::local::SharedData, namespace: String) -> Self {
        Self { storage, namespace }
    }

    fn convert_store_mode(
        &self,
        mode: &StoreMode,
    ) -> Result<crate::local::StoreMode, DataStorageError> {
        match mode {
            StoreMode::Always => Ok(crate::local::StoreMode::Always),
            StoreMode::Absent => Ok(crate::local::StoreMode::Absent),
            StoreMode::Cas(cas_str) => {
                let cas: u32 = cas_str.parse()?;
                Ok(crate::local::StoreMode::Cas(cas))
            }
        }
    }

    fn namespaced_key(&self, key: &str) -> String {
        format!("{}:{}", self.namespace, key)
    }
}

impl DataStorage for LocalDataStorage {
    async fn get_keys(&self) -> Result<Vec<String>, DataStorageError> {
        let all_keys = self.storage.keys();
        let namespace_prefix = format!("{}:", self.namespace);

        // Filter keys to only include those that belong to this namespace
        let filtered_keys: Vec<String> = all_keys
            .into_iter()
            .filter(|key| key.starts_with(&namespace_prefix))
            .map(|key| {
                // Remove the namespace prefix to return just the key name
                key.strip_prefix(&namespace_prefix)
                    .unwrap_or(&key)
                    .to_string()
            })
            .collect();

        Ok(filtered_keys)
    }

    async fn store<T: Serialize>(
        &self,
        key: &str,
        mode: &StoreMode,
        item: &T,
    ) -> Result<(), DataStorageError> {
        let serialized = bincode::serialize(item)?;
        let local_mode = self.convert_store_mode(mode)?;
        let namespaced_key = self.namespaced_key(key);
        self.storage.set(&namespaced_key, &serialized, local_mode)?;
        Ok(())
    }

    async fn get<T: DeserializeOwned>(
        &self,
        key: &str,
    ) -> Result<Option<(T, String)>, DataStorageError> {
        let namespaced_key = self.namespaced_key(key);
        match self.storage.get(&namespaced_key)? {
            Some((data, cas)) => {
                let deserialized: T =
                    bincode::deserialize(&data).map_err(DataStorageError::from)?;
                Ok(Some((deserialized, cas.to_string())))
            }
            None => Ok(None),
        }
    }

    async fn delete(&self, key: &str) -> Result<(), DataStorageError> {
        let namespaced_key = self.namespaced_key(key);
        self.storage.delete(&namespaced_key)?;
        Ok(())
    }

    async fn delete_all(&self) -> Result<(), DataStorageError> {
        let all_keys = self.storage.keys();
        let namespace_prefix = format!("{}:", self.namespace);

        // Only delete keys that belong to this namespace
        for key in all_keys {
            if key.starts_with(&namespace_prefix) {
                self.storage.delete(&key)?;
            }
        }
        Ok(())
    }
}

/// A distributed data storage implementation that stores data across multiple nodes.
pub struct RemoteDataStorage {
    storage: Rc<crate::distributed::DistributedStorageClient>,
    sanitized_store: String,
    sanitized_partition: String,
    ttl_millis: u32,
}

impl RemoteDataStorage {
    /// Creates a high level data storage instance for the given `store`/`partition`,
    /// using `storage` and a default TTL in milliseconds.
    pub(crate) fn new(
        storage: Rc<crate::distributed::DistributedStorageClient>,
        store: String,
        partition: String,
        ttl_millis: u32,
    ) -> Self {
        // Sanitize store and partition names using form_urlencoded
        let sanitized_store = form_urlencoded::byte_serialize(store.as_bytes()).collect();
        let sanitized_partition = form_urlencoded::byte_serialize(partition.as_bytes()).collect();
        Self {
            storage,
            sanitized_store,
            sanitized_partition,
            ttl_millis,
        }
    }

    fn convert_store_mode(&self, mode: &StoreMode) -> crate::distributed::StoreMode {
        match mode {
            StoreMode::Always => crate::distributed::StoreMode::Always,
            StoreMode::Absent => crate::distributed::StoreMode::Absent,
            StoreMode::Cas(cas_str) => crate::distributed::StoreMode::Cas(cas_str.clone()),
        }
    }

    fn sanitize_key(&self, key: &str) -> String {
        form_urlencoded::byte_serialize(key.as_bytes()).collect()
    }
}

impl DataStorage for RemoteDataStorage {
    async fn get_keys(&self) -> Result<Vec<String>, DataStorageError> {
        // Try to get keys, if store doesn't exist, return empty list
        match self
            .storage
            .get_keys(&self.sanitized_store, &self.sanitized_partition)
            .await
        {
            Ok(keys) => {
                // Decode the keys to return them in the same format as they were stored
                let decoded_keys: Vec<String> = keys
                    .into_iter()
                    .filter_map(|encoded_key| {
                        let decoded = form_urlencoded::parse(encoded_key.as_bytes())
                            .next()
                            .map(|(key, _)| key.into_owned());

                        if decoded.is_none() {
                            logger::debug!("Key not URL-encoded or decode failed: {encoded_key}");
                        }

                        decoded
                    })
                    .collect();

                Ok(decoded_keys)
            }
            Err(e) => {
                logger::warn!("Error getting keys: {e}");
                Ok(vec![])
            }
        }
    }

    async fn store<T: Serialize>(
        &self,
        key: &str,
        mode: &StoreMode,
        item: &T,
    ) -> Result<(), DataStorageError> {
        let serialized = bincode::serialize(item)?;
        let distributed_mode = self.convert_store_mode(mode);
        let sanitized_key = self.sanitize_key(key);

        // Try to store first
        match self
            .storage
            .store(
                &self.sanitized_store,
                &self.sanitized_partition,
                &sanitized_key,
                &distributed_mode,
                &serialized,
            )
            .await
        {
            Ok(()) => Ok(()),
            Err(crate::distributed::DistributedStorageError::StoreNotFound) => {
                // Store doesn't exist, create it and retry
                let store = crate::distributed::Store::new(
                    self.sanitized_store.clone(),
                    Some(self.ttl_millis),
                    None,
                );

                // Try to create the store, ignore if it already exists
                if let Err(e) = self.storage.upsert_store(&store).await {
                    logger::warn!("Error creating store: {e}");
                }

                // Try storing again
                self.storage
                    .store(
                        &self.sanitized_store,
                        &self.sanitized_partition,
                        &sanitized_key,
                        &distributed_mode,
                        &serialized,
                    )
                    .await?;
                Ok(())
            }
            Err(e) => Err(e.into()), // Other errors, propagate them
        }
    }

    async fn get<T: DeserializeOwned>(
        &self,
        key: &str,
    ) -> Result<Option<(T, String)>, DataStorageError> {
        let sanitized_key = self.sanitize_key(key);
        match self
            .storage
            .get(
                &self.sanitized_store,
                &self.sanitized_partition,
                &sanitized_key,
            )
            .await
        {
            Ok((data, cas)) => {
                let deserialized: T =
                    bincode::deserialize(&data).map_err(DataStorageError::from)?;
                Ok(Some((deserialized, cas)))
            }
            Err(crate::distributed::DistributedStorageError::KeyNotFound) => {
                logger::debug!("Key not found: {key}");
                Ok(None)
            }
            Err(e) => {
                logger::error!("Error getting value for key {key}: {e:?}");
                Err(e.into())
            }
        }
    }

    async fn delete(&self, key: &str) -> Result<(), DataStorageError> {
        // Try to delete, ignore if store doesn't exist
        let sanitized_key = self.sanitize_key(key);
        if let Err(e) = self
            .storage
            .delete(
                &self.sanitized_store,
                &self.sanitized_partition,
                &sanitized_key,
            )
            .await
        {
            logger::warn!("Error deleting key {key}: {e}");
        }
        Ok(())
    }

    async fn delete_all(&self) -> Result<(), DataStorageError> {
        // Try to delete partition, ignore if store doesn't exist
        if let Err(e) = self
            .storage
            .delete_partition(&self.sanitized_store, &self.sanitized_partition)
            .await
        {
            logger::warn!("Error deleting partition: {e}");
        }
        Ok(())
    }
}

/// Builder for creating data storage instances.
///
/// Provides methods to create local and distributed storage instances with
/// configurable settings. The builder pattern allows for flexible configuration
/// of storage behavior.
///
/// # Examples
///
/// ```rust
/// # use data_storage_lib::{DataStorageBuilder, DataStorage};
/// # async fn example(builder: DataStorageBuilder) {
/// // Create a local storage instance
/// let local_storage = builder.local("my-local-storage");
///
/// // Create a distributed storage instance with 60-second TTL
/// let remote_storage = builder.remote("my-remote-storage", 60000);
/// # }
/// ```
pub struct DataStorageBuilder {
    prefix: String,
    shared_data: Rc<crate::local::SharedData>,
    distributed_storage: Option<Rc<crate::distributed::DistributedStorageClient>>,
}
/// DataStorageBuilder can be injected in your configuration function.
/// ```rust
/// #[entrypoint]
/// async fn configure(
///     launcher: Launcher,
///     store_builder: DataStorageBuilder,
///     Configuration(configuration): Configuration,
/// ) -> anyhow::Result<()> {
/// }
/// ```
impl FromContext<ConfigureContext> for DataStorageBuilder {
    type Error = DataStorageBuilderError;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        // Extract local storage (required)
        let shared_data: crate::local::SharedData = context
            .extract()
            .map_err(|_| DataStorageBuilderError::LocalStorageRequired)?;
        // Extract distributed storage (optional - will be None if not available)
        let distributed_storage: Result<crate::distributed::DistributedStorageClient, _> =
            context.extract();
        // Extract metadata for policy isolation
        let metadata: pdk_core::policy_context::api::Metadata = context
            .extract()
            .map_err(|_| DataStorageBuilderError::MetadataRequired)?;

        let prefix = format!(
            "isolated-storage-{}-{}",
            metadata.policy_metadata.policy_name, metadata.policy_metadata.policy_namespace
        );

        pdk_core::logger::info!(
            "DataStorageBuilder: creating prefix '{}' for policy '{}' in namespace '{}'",
            prefix,
            metadata.policy_metadata.policy_name,
            metadata.policy_metadata.policy_namespace
        );

        Ok(DataStorageBuilder {
            prefix,
            shared_data: Rc::new(shared_data),
            distributed_storage: distributed_storage.ok().map(Rc::new),
        })
    }
}

impl DataStorageBuilder {
    /// Indicates that storage should not be isolated to a single policy.
    /// The resulting state will be shared across policy instances that use
    /// the same storage ID.
    pub fn shared(mut self) -> Self {
        self.prefix = "shared-storage".to_string();
        self
    }

    /// Creates a local data storage instance identified by the key.
    pub fn local<T: Into<String>>(&self, key: T) -> LocalDataStorage {
        let key_str = key.into();
        // Create a truly unique namespace by combining prefix and key
        // This ensures each policy gets its own isolated storage
        let namespace = format!("{}-{}", self.prefix, key_str);

        pdk_core::logger::info!(
            "DataStorageBuilder::local: creating namespace '{}' with prefix '{}' and key '{}'",
            namespace,
            self.prefix,
            key_str
        );

        LocalDataStorage::new((*self.shared_data).clone(), namespace)
    }

    /// Creates a distributed data storage instance identified by the provided key and ttl in milliseconds.
    ///
    /// **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.
    ///
    /// # Panics
    ///
    /// Panics if distributed storage is not available in the current context.
    /// Ensure that distributed storage is properly configured before calling this method.
    pub fn remote<T: Into<String>>(&self, key: T, ttl_millis: u32) -> RemoteDataStorage {
        let key_str = key.into();
        let storage = self
            .distributed_storage
            .as_ref()
            .expect("Distributed storage not available - check if it's configured");
        RemoteDataStorage::new(Rc::clone(storage), key_str.clone(), key_str, ttl_millis)
    }
}