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