Skip to main content

unb_runtime/
discover.rs

1use std::collections::BTreeSet;
2use std::pin::Pin;
3
4use futures_util::stream::{unfold, Stream};
5use unb_core::{Detail, DiscoverEvent, DiscoverPlan, Mode, Scope, DEFAULT_HOPS};
6
7use crate::error::WsError;
8use crate::wire::Wire;
9
10impl Wire {
11    /// Open a `Discover` walk and stream its typed catalog events.
12    ///
13    /// Yields one [`DiscoverEvent`] per graph observation — `NodeCatalog`, `Edge`,
14    /// `Warning` — and finally the `Done` marker (carrying the `discover_id`),
15    /// after which the stream ends. An error terminal or a closed session ends the
16    /// stream without a `Done`.
17    pub async fn discover_catalog(
18        &self,
19        target_path: &str,
20        detail: Detail,
21        scope: Scope,
22    ) -> Result<Pin<Box<dyn Stream<Item = DiscoverEvent> + '_>>, WsError> {
23        let plan = DiscoverPlan {
24            discover_id: String::new(),
25            detail,
26            scope,
27            hops: DEFAULT_HOPS,
28            visited: BTreeSet::new(),
29            timeout_ms: None,
30            mode: Mode::PartialOk,
31        };
32        let stream = self.client_session().discover(target_path, plan).await?;
33
34        Ok(Box::pin(unfold(
35            (stream, false),
36            |(mut stream, done)| async move {
37                if done {
38                    return None;
39                }
40                while let Ok(Some(envelope)) = stream.next().await {
41                    match serde_json::from_slice::<DiscoverEvent>(&envelope.payload) {
42                        Ok(marker @ DiscoverEvent::Done { .. }) => {
43                            return Some((marker, (stream, true)))
44                        }
45                        Ok(event) => return Some((event, (stream, false))),
46                        Err(_) => continue,
47                    }
48                }
49                None
50            },
51        )))
52    }
53}