hashtree_cli/nostrdb_integration/
mod.rs1pub mod access;
4pub mod crawler;
5pub mod snapshot;
6
7pub use nostrdb_social::Ndb;
8use nostrdb_social::{Config as NdbConfig, Transaction};
9use std::collections::HashMap;
10use std::path::Path;
11use std::sync::Arc;
12
13#[cfg(test)]
14use std::sync::{Mutex, MutexGuard, OnceLock};
15
16#[cfg(test)]
17pub type TestLockGuard = MutexGuard<'static, ()>;
18
19#[cfg(test)]
20static NDB_TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
21
22#[cfg(test)]
23pub fn test_lock() -> TestLockGuard {
24 NDB_TEST_LOCK.get_or_init(|| Mutex::new(())).lock().unwrap()
25}
26
27pub use access::SocialGraphAccessControl;
28pub use crawler::SocialGraphCrawler;
29
30#[derive(Debug, Clone, Default, serde::Serialize)]
32pub struct SocialGraphStats {
33 pub root: Option<String>,
34 pub total_follows: usize,
35 pub max_depth: u32,
36 pub enabled: bool,
37}
38
39pub fn init_ndb(data_dir: &Path) -> anyhow::Result<Arc<Ndb>> {
41 init_ndb_with_mapsize(data_dir, None)
42}
43
44pub fn init_ndb_with_mapsize(
46 data_dir: &Path,
47 mapsize_bytes: Option<u64>,
48) -> anyhow::Result<Arc<Ndb>> {
49 let ndb_dir = data_dir.join("nostrdb_social");
50 init_ndb_at_path(&ndb_dir, mapsize_bytes)
51}
52
53pub fn init_ndb_at_path(db_dir: &Path, mapsize_bytes: Option<u64>) -> anyhow::Result<Arc<Ndb>> {
55 std::fs::create_dir_all(db_dir)?;
56 let mut config = NdbConfig::new().set_ingester_threads(2);
57 if let Some(bytes) = mapsize_bytes {
58 let mapsize = usize::try_from(bytes).unwrap_or(usize::MAX);
59 config = config.set_mapsize(mapsize);
60 }
61 let ndb = Ndb::new(db_dir.to_str().unwrap_or("."), &config)?;
62 Ok(Arc::new(ndb))
63}
64
65pub fn set_social_graph_root(ndb: &Ndb, pk_bytes: &[u8; 32]) {
67 nostrdb_social::socialgraph::set_root(ndb, pk_bytes);
68}
69
70pub fn get_follow_distance(ndb: &Ndb, pk_bytes: &[u8; 32]) -> Option<u32> {
73 let txn = Transaction::new(ndb).ok()?;
74 let distance = nostrdb_social::socialgraph::get_follow_distance(&txn, ndb, pk_bytes);
75 if distance >= 1000 {
76 None
77 } else {
78 Some(distance)
79 }
80}
81
82pub fn get_follows(ndb: &Ndb, pk_bytes: &[u8; 32]) -> Vec<[u8; 32]> {
84 let txn = match Transaction::new(ndb) {
85 Ok(t) => t,
86 Err(_) => return Vec::new(),
87 };
88 nostrdb_social::socialgraph::get_followed(&txn, ndb, pk_bytes, 10000)
89}
90
91fn clamp_socialgraph_list(count: usize) -> usize {
92 let max = i32::MAX as usize;
93 if count > max {
94 max
95 } else {
96 count
97 }
98}
99
100pub fn is_overmuted(ndb: &Ndb, root_pk: &[u8; 32], user_pk: &[u8; 32], threshold: f64) -> bool {
103 if threshold <= 0.0 {
104 return false;
105 }
106 if user_pk == root_pk {
107 return false;
108 }
109
110 let txn = match Transaction::new(ndb) {
111 Ok(t) => t,
112 Err(_) => return false,
113 };
114
115 let muter_count = nostrdb_social::socialgraph::muter_count(&txn, ndb, user_pk);
116 if muter_count == 0 {
117 return false;
118 }
119
120 if nostrdb_social::socialgraph::is_muting(&txn, ndb, root_pk, user_pk) {
121 return true;
122 }
123
124 let follower_count = nostrdb_social::socialgraph::follower_count(&txn, ndb, user_pk);
125
126 let mut stats: HashMap<u32, (usize, usize)> = HashMap::new();
127
128 let followers = nostrdb_social::socialgraph::get_followers(
129 &txn,
130 ndb,
131 user_pk,
132 clamp_socialgraph_list(follower_count),
133 );
134 for follower_pk in followers {
135 let distance = nostrdb_social::socialgraph::get_follow_distance(&txn, ndb, &follower_pk);
136 if distance >= 1000 {
137 continue;
138 }
139 let entry = stats.entry(distance).or_insert((0, 0));
140 entry.0 += 1;
141 }
142
143 let muters = nostrdb_social::socialgraph::get_muters(
144 &txn,
145 ndb,
146 user_pk,
147 clamp_socialgraph_list(muter_count),
148 );
149 for muter_pk in muters {
150 let distance = nostrdb_social::socialgraph::get_follow_distance(&txn, ndb, &muter_pk);
151 if distance >= 1000 {
152 continue;
153 }
154 let entry = stats.entry(distance).or_insert((0, 0));
155 entry.1 += 1;
156 }
157
158 let mut distances: Vec<u32> = stats.keys().cloned().collect();
159 distances.sort_unstable();
160
161 for distance in distances {
162 let (followers, muters) = stats[&distance];
163 if followers + muters > 0 {
164 return (muters as f64) * threshold > (followers as f64);
165 }
166 }
167
168 false
169}
170
171pub fn ingest_event(ndb: &Ndb, sub_id: &str, event_json: &str) {
174 let relay_msg = format!(r#"["EVENT","{}",{}]"#, sub_id, event_json);
175 if let Err(e) = ndb.process_event(&relay_msg) {
176 tracing::warn!("Failed to ingest event into nostrdb_social: {}", e);
177 }
178}
179
180#[cfg(test)]
181mod tests {
182 use super::*;
183 use tempfile::TempDir;
184
185 #[test]
186 fn test_init_ndb() {
187 let _guard = test_lock();
188 let tmp = TempDir::new().unwrap();
189 let ndb = init_ndb(tmp.path()).unwrap();
190 assert!(Arc::strong_count(&ndb) == 1);
191 }
192
193 #[test]
194 fn test_set_root_and_get_follow_distance() {
195 let _guard = test_lock();
196 let tmp = TempDir::new().unwrap();
197 let ndb = init_ndb(tmp.path()).unwrap();
198 let root_pk = [1u8; 32];
199 set_social_graph_root(&ndb, &root_pk);
200 std::thread::sleep(std::time::Duration::from_millis(100));
202 let dist = get_follow_distance(&ndb, &root_pk);
203 assert_eq!(dist, Some(0));
204 }
205
206 #[test]
207 fn test_unknown_pubkey_follow_distance() {
208 let _guard = test_lock();
209 let tmp = TempDir::new().unwrap();
210 let ndb = init_ndb(tmp.path()).unwrap();
211 let root_pk = [1u8; 32];
212 set_social_graph_root(&ndb, &root_pk);
213 std::thread::sleep(std::time::Duration::from_millis(100));
214 let unknown_pk = [2u8; 32];
215 assert_eq!(get_follow_distance(&ndb, &unknown_pk), None);
216 }
217
218 #[test]
219 fn test_ingest_event_no_panic() {
220 let _guard = test_lock();
221 let tmp = TempDir::new().unwrap();
222 let ndb = init_ndb(tmp.path()).unwrap();
223 ingest_event(&ndb, "sub1", r#"{"kind":1,"content":"hello"}"#);
225 }
226
227 #[test]
228 fn test_get_follows_empty() {
229 let _guard = test_lock();
230 let tmp = TempDir::new().unwrap();
231 let ndb = init_ndb(tmp.path()).unwrap();
232 let pk = [1u8; 32];
233 assert!(get_follows(&ndb, &pk).is_empty());
234 }
235}