Skip to main content

faucet_cli/
state.rs

1//! Build a `StateStore` from a `(kind, config)` pair, with feature-gated
2//! backends for Redis and PostgreSQL.
3
4use crate::config::StateStoreSpec;
5use crate::error::{CliError, CliResult};
6use faucet_core::{FileStateStore, MemoryStateStore, StateStore};
7use serde::Deserialize;
8use std::path::PathBuf;
9use std::sync::Arc;
10
11/// Config block for the built-in file backend.
12#[derive(Debug, Deserialize)]
13struct FileStateConfig {
14    path: PathBuf,
15}
16
17#[cfg(feature = "state-redis")]
18#[derive(Debug, Deserialize)]
19struct RedisStateConfig {
20    url: String,
21    #[serde(default = "default_redis_namespace")]
22    namespace: String,
23}
24
25#[cfg(feature = "state-redis")]
26fn default_redis_namespace() -> String {
27    "faucet".to_owned()
28}
29
30#[cfg(feature = "state-postgres")]
31#[derive(Debug, Deserialize)]
32struct PostgresStateConfig {
33    url: String,
34    #[serde(default = "default_pg_table")]
35    table: String,
36    #[serde(default)]
37    ensure_table: bool,
38}
39
40#[cfg(feature = "state-postgres")]
41fn default_pg_table() -> String {
42    "faucet_state".to_owned()
43}
44
45/// Construct a state store from the parsed `state:` block.
46pub async fn build_state_store(spec: &StateStoreSpec) -> CliResult<Arc<dyn StateStore>> {
47    match spec.kind.as_str() {
48        "memory" => Ok(Arc::new(MemoryStateStore::new())),
49        "file" => {
50            let cfg = decode::<FileStateConfig>("file", spec.config.clone())?;
51            Ok(Arc::new(FileStateStore::new(cfg.path)))
52        }
53        #[cfg(feature = "state-redis")]
54        "redis" => {
55            let cfg = decode::<RedisStateConfig>("redis", spec.config.clone())?;
56            Ok(Arc::new(
57                faucet_state_redis::RedisStateStore::connect(&cfg.url, &cfg.namespace).await?,
58            ))
59        }
60        #[cfg(feature = "state-postgres")]
61        "postgres" => {
62            let cfg = decode::<PostgresStateConfig>("postgres", spec.config.clone())?;
63            let store =
64                faucet_state_postgres::PostgresStateStore::connect_with(&cfg.url, 5, &cfg.table)
65                    .await?;
66            if cfg.ensure_table {
67                store.ensure_table().await?;
68            }
69            Ok(Arc::new(store))
70        }
71        other => Err(CliError::UnknownStateStore {
72            name: other.to_owned(),
73            available: available_state_kinds().join(", "),
74        }),
75    }
76}
77
78/// Names of every state-store backend compiled into this build.
79pub fn available_state_kinds() -> Vec<&'static str> {
80    let mut v = vec!["memory", "file"];
81    #[cfg(feature = "state-redis")]
82    v.push("redis");
83    #[cfg(feature = "state-postgres")]
84    v.push("postgres");
85    v
86}
87
88fn decode<T: serde::de::DeserializeOwned>(name: &str, config: serde_json::Value) -> CliResult<T> {
89    serde_json::from_value(config).map_err(|e| CliError::InvalidConnectorConfig {
90        kind: "state",
91        name: name.to_owned(),
92        message: e.to_string(),
93    })
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use serde_json::json;
100
101    #[tokio::test]
102    async fn builds_memory_store() {
103        let spec = StateStoreSpec {
104            kind: "memory".into(),
105            config: json!({}),
106        };
107        let store = build_state_store(&spec).await.unwrap();
108        store.put("k", &json!(1)).await.unwrap();
109        assert_eq!(store.get("k").await.unwrap(), Some(json!(1)));
110    }
111
112    #[tokio::test]
113    async fn builds_file_store() {
114        let dir = tempfile::tempdir().unwrap();
115        let spec = StateStoreSpec {
116            kind: "file".into(),
117            config: json!({"path": dir.path().to_str().unwrap()}),
118        };
119        let store = build_state_store(&spec).await.unwrap();
120        store.put("k", &json!("v")).await.unwrap();
121    }
122
123    #[tokio::test]
124    async fn unknown_kind_errors() {
125        let spec = StateStoreSpec {
126            kind: "nope".into(),
127            config: json!({}),
128        };
129        let err = build_state_store(&spec).await.err().expect("should fail");
130        match err {
131            CliError::UnknownStateStore { name, .. } => assert_eq!(name, "nope"),
132            other => panic!("expected UnknownStateStore, got {other:?}"),
133        }
134    }
135}