1use 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#[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
45pub 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
78pub 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}