Skip to main content

statsig_rust/
data_store_interface.rs

1use std::fmt::Display;
2
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5
6use crate::{hashing::HashUtil, StatsigErr, StatsigOptions};
7
8pub enum RequestPath {
9    RulesetsV2,
10    RulesetsV1,
11    IDListsV1,
12    IDList,
13}
14
15impl Display for RequestPath {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        let value = match self {
18            RequestPath::IDListsV1 => "/v1/get_id_lists",
19            RequestPath::IDList => "id_list",
20            RequestPath::RulesetsV2 => "/v2/download_config_specs",
21            RequestPath::RulesetsV1 => "/v1/download_config_specs",
22        };
23        write!(f, "{value}")
24    }
25}
26
27pub enum CompressFormat {
28    PlainText,
29    Gzip,
30    StatsigBr,
31}
32
33impl Display for CompressFormat {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        let value = match self {
36            CompressFormat::PlainText => "plain_text",
37            CompressFormat::Gzip => "gzip",
38            CompressFormat::StatsigBr => "statsig-br",
39        };
40        write!(f, "{value}")
41    }
42}
43
44#[derive(Deserialize, Serialize)]
45pub struct DataStoreResponse {
46    pub result: Option<String>,
47    pub time: Option<u64>,
48}
49
50#[derive(Deserialize, Serialize)]
51pub struct DataStoreBytesResponse {
52    pub result: Option<Vec<u8>>,
53    pub time: Option<u64>,
54}
55
56#[async_trait]
57pub trait DataStoreTrait: Send + Sync {
58    async fn initialize(&self) -> Result<(), StatsigErr>;
59    async fn shutdown(&self) -> Result<(), StatsigErr>;
60    async fn get(&self, key: &str) -> Result<DataStoreResponse, StatsigErr>;
61    async fn set(&self, key: &str, value: &str, time: Option<u64>) -> Result<(), StatsigErr>;
62    async fn set_bytes(
63        &self,
64        key: &str,
65        value: &[u8],
66        time: Option<u64>,
67    ) -> Result<(), StatsigErr> {
68        let value = std::str::from_utf8(value).map_err(|e| {
69            StatsigErr::DataStoreFailure(format!("Failed to decode bytes as UTF-8: {e}"))
70        })?;
71        self.set(key, value, time).await
72    }
73
74    async fn get_bytes(&self, key: &str) -> Result<DataStoreBytesResponse, StatsigErr> {
75        let response = self.get(key).await?;
76        Ok(DataStoreBytesResponse {
77            result: response.result.map(|value| value.into_bytes()),
78            time: response.time,
79        })
80    }
81
82    /// Whether this store can natively read/write bytes without UTF-8 conversion.
83    /// Defaults to `false` to preserve the existing string-first behavior.
84    fn supports_bytes(&self) -> bool {
85        false
86    }
87
88    async fn support_polling_updates_for(&self, path: RequestPath) -> bool;
89}
90
91#[derive(Clone, Debug, Default)]
92pub enum DataStoreKeyVersion {
93    #[default]
94    V2Hashed,
95    V3HumanReadable,
96}
97
98impl From<&str> for DataStoreKeyVersion {
99    fn from(level: &str) -> Self {
100        match level.to_lowercase().as_str() {
101            "v2" | "2" => DataStoreKeyVersion::V2Hashed,
102            "v3" | "3" => DataStoreKeyVersion::V3HumanReadable,
103            _ => DataStoreKeyVersion::default(),
104        }
105    }
106}
107
108#[must_use]
109pub(crate) fn get_data_store_key(
110    path: RequestPath,
111    sdk_key: &str,
112    hashing: &HashUtil,
113    options: &StatsigOptions,
114) -> String {
115    let compress_format = if options
116        .data_store
117        .as_ref()
118        .is_some_and(|data_store| data_store.supports_bytes())
119    {
120        CompressFormat::StatsigBr
121    } else {
122        CompressFormat::PlainText
123    };
124
125    let key = match options
126        .data_store_key_schema_version
127        .clone()
128        .unwrap_or_default()
129    {
130        DataStoreKeyVersion::V3HumanReadable => {
131            let mut key = sdk_key.to_string();
132            key.truncate(20);
133            key
134        }
135        DataStoreKeyVersion::V2Hashed => hashing.hash(sdk_key, &crate::HashAlgorithm::Sha256),
136    };
137
138    format!("statsig|{path}|{compress_format}|{key}")
139}