pdk-data-storage-lib 1.7.0-alpha.0

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

//! Local shared data storage primitives
//!
//! Provides a lightweight, in-memory shared data API used by policies within a
//! single node. This module defines the low-level [`LocalStorage`] trait and the
//! injectable [`SharedData`] implementation that integrates with the host
//! runtime.

use pdk_core::classy::extract::context::ConfigureContext;
use pdk_core::classy::extract::{Extract, FromContext};
use pdk_core::classy::proxy_wasm::types::Status;
use std::convert::Infallible;
use std::rc::Rc;
use thiserror::Error;

/// Defines the behavior for store operations.
#[derive(PartialEq, Eq, Debug, Copy, Clone)]
pub enum StoreMode {
    /// Indicates that the store operation does not depend on any condition.
    Always,
    /// Indicates that the store operation should succeed only if no other value was present in the store.
    Absent,
    /// Indicates that the store operation should succeed only if the stored value matches the provided cas.
    Cas(u32),
}

#[non_exhaustive]
#[derive(Error, Debug)]
/// Error type for local storage operations.
pub enum LocalStorageError {
    /// Indicates that the operation was aborted due to difference in the expected CAS value
    #[error("Cas mismatch.")]
    CasMismatch,
    /// Unhandled error from the proxy wasm runtime.
    #[error("Low level error.")]
    ProxyWasm(Status),
}

impl From<Status> for LocalStorageError {
    fn from(value: Status) -> Self {
        Self::ProxyWasm(value)
    }
}

/// The [`LocalStorage`] trait defines the methods for storing and retrieving data.
pub trait LocalStorage {
    /// Sets the value of a key in the store with the particular [`StoreMode`].
    fn set(&self, key: &str, value: &[u8], mode: StoreMode) -> Result<(), LocalStorageError>;

    /// Retrieves the value of a key from the store.
    fn get(&self, key: &str) -> Result<Option<(Vec<u8>, u32)>, LocalStorageError>;

    /// Deletes a key from the store.
    // proxy_wasm_rust_sdk is panicking when a cas mismatch is returned, therefore we cannot
    // take advantage of the cas for delete.
    fn delete(&self, key: &str) -> Result<(), LocalStorageError>;

    /// Returns all the keys in the store.
    fn keys(&self) -> Vec<String>;
}

#[derive(Clone)]
/// Implementation of [`LocalStorage`] trait for the shared data.
pub struct SharedData {
    shared_data: Rc<dyn pdk_core::classy::SharedData>,
}

impl SharedData {
    fn store_mode_to_cas(
        &self,
        key: &str,
        mode: StoreMode,
    ) -> Result<Option<u32>, LocalStorageError> {
        match mode {
            StoreMode::Always => Ok(None),
            StoreMode::Absent => {
                match self.shared_data.shared_data_get(key) {
                    (Some(data), cas) => {
                        if data.is_empty() {
                            Ok(Some(cas.unwrap_or(u32::MAX))) // Should never be undefined since it has data
                        } else {
                            Err(LocalStorageError::CasMismatch)
                        }
                    }
                    (None, Some(cas)) => Ok(Some(cas)), // Scenario of store with empty value
                    (None, None) => Ok(Some(u32::MAX)),
                }
            }
            StoreMode::Cas(cas) => Ok(Some(cas)),
        }
    }
}

impl FromContext<ConfigureContext> for SharedData {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        let shared_data: Rc<dyn pdk_core::classy::SharedData> = context.extract()?;
        Ok(Self { shared_data })
    }
}

impl LocalStorage for SharedData {
    fn set(&self, key: &str, value: &[u8], mode: StoreMode) -> Result<(), LocalStorageError> {
        let cas = self.store_mode_to_cas(key, mode)?;

        self.shared_data
            .shared_data_set(key, value, cas)
            .map_err(|e| match e {
                Status::CasMismatch => LocalStorageError::CasMismatch,
                v => LocalStorageError::ProxyWasm(v),
            })
    }

    fn get(&self, key: &str) -> Result<Option<(Vec<u8>, u32)>, LocalStorageError> {
        match self.shared_data.shared_data_get(key) {
            (Some(data), Some(cas)) => Ok(Some((data, cas))),
            _ => Ok(None),
        }
    }

    fn delete(&self, key: &str) -> Result<(), LocalStorageError> {
        let cas = self.store_mode_to_cas(key, StoreMode::Always)?;
        let _ = self.shared_data.shared_data_remove(key, cas)?;
        Ok(())
    }

    fn keys(&self) -> Vec<String> {
        self.shared_data.shared_data_keys()
    }
}