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    pub checksum: Option<String>,
49    pub has_updates: Option<bool>,
50}
51
52#[derive(Deserialize, Serialize)]
53pub struct DataStoreBytesResponse {
54    pub result: Option<Vec<u8>>,
55    pub time: Option<u64>,
56    pub checksum: Option<String>,
57    pub has_updates: Option<bool>,
58}
59
60#[derive(Clone, Debug)]
61pub struct DataStoreGetBytesRequest {
62    pub since_time: Option<u64>,
63    pub checksum: Option<String>,
64}
65
66#[derive(Clone, Debug)]
67pub(crate) struct DataStoreCacheKeys {
68    pub plain_text: String,
69    pub statsig_br: String,
70}
71
72impl DataStoreCacheKeys {
73    pub(crate) fn from_selected_key(key: &str) -> Self {
74        Self {
75            plain_text: data_store_key_with_compress_format(key, CompressFormat::PlainText),
76            statsig_br: data_store_key_with_compress_format(key, CompressFormat::StatsigBr),
77        }
78    }
79}
80
81#[async_trait]
82pub trait DataStoreTrait: Send + Sync {
83    async fn initialize(&self) -> Result<(), StatsigErr>;
84    async fn shutdown(&self) -> Result<(), StatsigErr>;
85    async fn get(&self, key: &str) -> Result<DataStoreResponse, StatsigErr>;
86    async fn set(&self, key: &str, value: &str, time: Option<u64>) -> Result<(), StatsigErr>;
87    async fn set_bytes(
88        &self,
89        key: &str,
90        value: &[u8],
91        time: Option<u64>,
92        checksum: Option<String>,
93    ) -> Result<(), StatsigErr> {
94        let _ = (key, value, time, checksum);
95        Err(StatsigErr::BytesNotImplemented)
96    }
97
98    async fn get_bytes(
99        &self,
100        _key: &str,
101        request: DataStoreGetBytesRequest,
102    ) -> Result<DataStoreBytesResponse, StatsigErr> {
103        let _ = request;
104        Err(StatsigErr::BytesNotImplemented)
105    }
106
107    async fn support_polling_updates_for(&self, path: RequestPath) -> bool;
108}
109
110#[derive(Clone, Debug, Default)]
111pub enum DataStoreKeyVersion {
112    #[default]
113    V2Hashed,
114    V3HumanReadable,
115}
116
117impl From<&str> for DataStoreKeyVersion {
118    fn from(level: &str) -> Self {
119        match level.to_lowercase().as_str() {
120            "v2" | "2" => DataStoreKeyVersion::V2Hashed,
121            "v3" | "3" => DataStoreKeyVersion::V3HumanReadable,
122            _ => DataStoreKeyVersion::default(),
123        }
124    }
125}
126
127#[must_use]
128pub(crate) fn get_data_store_key(
129    path: RequestPath,
130    sdk_key: &str,
131    hashing: &HashUtil,
132    options: &StatsigOptions,
133) -> String {
134    let key = match options
135        .data_store_key_schema_version
136        .clone()
137        .unwrap_or_default()
138    {
139        DataStoreKeyVersion::V3HumanReadable => {
140            let mut key = sdk_key.to_string();
141            key.truncate(20);
142            key
143        }
144        DataStoreKeyVersion::V2Hashed => hashing.hash(sdk_key, &crate::HashAlgorithm::Sha256),
145    };
146
147    format!("statsig|{path}|{}|{key}", CompressFormat::PlainText)
148}
149
150fn data_store_key_with_compress_format(key: &str, compress_format: CompressFormat) -> String {
151    let target = format!("|{compress_format}|");
152
153    [
154        CompressFormat::PlainText,
155        CompressFormat::Gzip,
156        CompressFormat::StatsigBr,
157    ]
158    .into_iter()
159    .map(|format| format!("|{format}|"))
160    .find_map(|current| {
161        key.contains(&current)
162            .then(|| key.replacen(&current, &target, 1))
163    })
164    .unwrap_or_else(|| key.to_string())
165}