Skip to main content

unb_runtime/
discover.rs

1use std::collections::BTreeSet;
2use std::pin::Pin;
3
4use unb_core::{Detail, DiscoverEvent, DiscoverPlan, Mode, Scope, DEFAULT_HOPS};
5use futures_util::stream::{unfold, Stream};
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        detail: Detail,
20        scope: Scope,
21    ) -> Result<Pin<Box<dyn Stream<Item = DiscoverEvent> + '_>>, WsError> {
22        let plan = DiscoverPlan {
23            discover_id: String::new(),
24            detail,
25            scope,
26            hops: DEFAULT_HOPS,
27            visited: BTreeSet::new(),
28            timeout_ms: None,
29            mode: Mode::PartialOk,
30        };
31        let stream = self.client_session().discover(plan).await?;
32
33        Ok(Box::pin(unfold(
34            (stream, false),
35            |(mut stream, done)| async move {
36                if done {
37                    return None;
38                }
39                while let Ok(Some(envelope)) = stream.next().await {
40                    match serde_json::from_slice::<DiscoverEvent>(&envelope.payload) {
41                        Ok(marker @ DiscoverEvent::Done { .. }) => {
42                            return Some((marker, (stream, true)))
43                        }
44                        Ok(event) => return Some((event, (stream, false))),
45                        Err(_) => continue,
46                    }
47                }
48                None
49            },
50        )))
51    }
52}