launchdarkly-server-sdk-redis 1.0.0-rc.1

use std::{collections::HashMap, convert::TryInto};

use launchdarkly_server_sdk::{
    AllData, DataKind, Flag, PersistentDataStore, PersistentStoreError, Segment, SerializedItem,
    StorageItem, Versioned,
};
use redis::{Client, Commands, Connection, ErrorKind, RedisError, Value};

pub struct RedisPersistentDataStore {
    client: Client,
    prefix: String,
}

impl RedisPersistentDataStore {
    pub(crate) fn new(client: Client, prefix: String) -> Self {
        Self { client, prefix }
    }

    fn get_connection(&self) -> Result<Connection, PersistentStoreError> {
        self.client
            .get_connection()
            .map_err(|e| PersistentStoreError::new(e.to_string()))
    }

    fn inited_key(&self) -> String {
        format!("{}:$inited", self.prefix)
    }

    fn flag_cache_key(&self) -> String {
        format!("{}:features", self.prefix)
    }

    fn segment_cache_key(&self) -> String {
        format!("{}:segments", self.prefix)
    }

    fn upsert_flag(
        &self,
        key: &str,
        serialized_item: SerializedItem,
    ) -> Result<bool, PersistentStoreError> {
        let mut connection = self.get_connection()?;
        redis::transaction(&mut connection, &[self.flag_cache_key()], |conn, pipe| {
            let result: Result<Option<String>, RedisError> = conn.hget(self.flag_cache_key(), key);
            match result {
                Ok(Some(json_string)) => {
                    let existing_item = SerializedItem {
                        version: 0,
                        deleted: false,
                        serialized_item: json_string,
                    };
                    let result: Result<StorageItem<Flag>, serde_json::Error> =
                        existing_item.try_into();
                    // If deserialization fails, then the data in redis is unusable so we can safely
                    // overwrite it.
                    if let Ok(flag) = result {
                        if flag.is_greater_than_or_equal(serialized_item.version) {
                            return Ok(Some(false));
                        }
                    }
                }
                Ok(None) => (),
                Err(e) => {
                    return Err(RedisError::from((
                        ErrorKind::ResponseError,
                        "failed to retrieve flag",
                        e.to_string(),
                    )))
                }
            };

            let result: Result<Value, RedisError> = pipe
                .hset(
                    self.flag_cache_key(),
                    key,
                    serialized_item.serialized_item.clone(),
                )
                .query(conn);

            match result {
                // If the result of the pipe is nil, then we know the watched key was modified.
                // Returning None causes the transaction to try again.
                Ok(Value::Nil) => Ok(None),
                Ok(_) => Ok(Some(true)),
                Err(e) => Err(e),
            }
        })
        .map_err(|e| PersistentStoreError::new(e.to_string()))
    }

    fn upsert_segment(
        &self,
        key: &str,
        serialized_item: SerializedItem,
    ) -> Result<bool, PersistentStoreError> {
        let mut connection = self.get_connection()?;
        redis::transaction(
            &mut connection,
            &[self.segment_cache_key()],
            |conn, pipe| {
                let result: Result<Option<String>, RedisError> =
                    conn.hget(self.segment_cache_key(), key);
                match result {
                    Ok(Some(json_string)) => {
                        let existing_item = SerializedItem {
                            version: 0,
                            deleted: false,
                            serialized_item: json_string,
                        };
                        let result: Result<StorageItem<Segment>, serde_json::Error> =
                            existing_item.try_into();
                        // If serialization fails, then the data in redis is unusable so we can safely
                        // overwrite it.
                        if let Ok(segment) = result {
                            if segment.is_greater_than_or_equal(serialized_item.version) {
                                return Ok(Some(false));
                            }
                        }
                    }
                    Ok(None) => (),
                    Err(e) => {
                        return Err(RedisError::from((
                            ErrorKind::ResponseError,
                            "failed to retrieve segment",
                            e.to_string(),
                        )))
                    }
                };

                let result: Result<Value, RedisError> = pipe
                    .hset(
                        self.segment_cache_key(),
                        key,
                        serialized_item.serialized_item.clone(),
                    )
                    .query(conn);

                match result {
                    // If the result of the pipe is nil, then we know the watched key was modified.
                    // Returning None causes the transaction to try again.
                    Ok(Value::Nil) => Ok(None),
                    Ok(_) => Ok(Some(true)),
                    Err(e) => Err(e),
                }
            },
        )
        .map_err(|e| PersistentStoreError::new(e.to_string()))
    }
}

impl PersistentDataStore for RedisPersistentDataStore {
    fn init(
        &mut self,
        all_data: AllData<SerializedItem, SerializedItem>,
    ) -> Result<(), PersistentStoreError> {
        let mut connection = self.get_connection()?;

        let mut pipeline = redis::pipe();
        pipeline
            .atomic()
            .del(self.flag_cache_key())
            .del(self.segment_cache_key())
            .set(self.inited_key(), "");

        all_data.flags.into_iter().for_each(|(key, item)| {
            pipeline.hset(self.flag_cache_key(), key, item.serialized_item);
        });

        all_data.segments.into_iter().for_each(|(key, item)| {
            pipeline.hset(self.segment_cache_key(), key, item.serialized_item);
        });

        let _: Value = pipeline
            .query(&mut connection)
            .map_err(|e| PersistentStoreError::new(e.to_string()))?;

        Ok(())
    }

    fn flag(&self, key: &str) -> Result<Option<SerializedItem>, PersistentStoreError> {
        let mut connection = self.get_connection()?;

        let result: Option<String> = connection
            .hget(self.flag_cache_key(), key)
            .map_err(|e| PersistentStoreError::new(e.to_string()))?;

        match result {
            Some(json_string) => Ok(Some(SerializedItem {
                version: 0,
                deleted: false,
                serialized_item: json_string,
            })),
            None => Ok(None),
        }
    }

    fn segment(&self, key: &str) -> Result<Option<SerializedItem>, PersistentStoreError> {
        let mut connection = self.get_connection()?;

        let result: Option<String> = connection
            .hget(self.segment_cache_key(), key)
            .map_err(|e| PersistentStoreError::new(e.to_string()))?;

        match result {
            Some(json_string) => Ok(Some(SerializedItem {
                version: 0,
                deleted: false,
                serialized_item: json_string,
            })),
            None => Ok(None),
        }
    }

    fn all_flags(&self) -> Result<HashMap<String, SerializedItem>, PersistentStoreError> {
        let mut connection = self.get_connection()?;
        let map: HashMap<String, String> = connection.hgetall(self.flag_cache_key()).expect("FAIL");

        Ok(map
            .into_iter()
            .map(|(key, serialized_item)| {
                (
                    key,
                    SerializedItem {
                        version: 0,
                        deleted: false,
                        serialized_item,
                    },
                )
            })
            .collect())
    }

    fn upsert(
        &mut self,
        kind: DataKind,
        key: &str,
        serialized_item: SerializedItem,
    ) -> Result<bool, PersistentStoreError> {
        match kind {
            DataKind::Flag => self.upsert_flag(key, serialized_item),
            DataKind::Segment => self.upsert_segment(key, serialized_item),
        }
    }

    fn is_initialized(&self) -> bool {
        if let Ok(mut connection) = self.get_connection() {
            let result: Result<bool, RedisError> = connection.exists(self.inited_key());

            match result {
                Ok(r) => r,
                _ => false,
            }
        } else {
            false
        }
    }
}

#[cfg(test)]
mod tests {
    use super::RedisPersistentDataStore;

    #[test]
    fn setting_prefix_affects_keys_correctly() {
        let client = redis::Client::open("redis://localhost:6379").expect("client created");
        let store = RedisPersistentDataStore::new(client, String::from("my-prefix"));

        assert_eq!("my-prefix:features", store.flag_cache_key());
        assert_eq!("my-prefix:segments", store.segment_cache_key());
        assert_eq!("my-prefix:$inited", store.inited_key());
    }
}