stac_io/
stream.rs

1//! Streaming JSON writer for a search-response [`stac::api::ItemCollection`]: the
2//! `features` array is written one item at a time, and the rest of the collection (links, context,
3//! counts) is supplied by a `finalize` callback after the items drain (the `next` link needs the last
4//! item; a `numberMatched` count may run concurrently).
5
6use crate::Result;
7use futures::{Stream, StreamExt};
8use serde_json::Value;
9use stac::api::{ItemCollection, Search};
10use std::{future::Future, io::Write, pin::Pin};
11
12/// A boxed, pinned stream of serialized STAC items.
13pub type ItemStream = Pin<Box<dyn Stream<Item = Result<Value>> + Send>>;
14
15/// Produces the finished [`ItemCollection`] (with empty `items`; the writer fills `numberReturned`) from
16/// the first item, the last item, and the number written. Called once, after the stream drains.
17pub type Finalize = Box<
18    dyn FnOnce(
19            Option<Value>,
20            Option<Value>,
21            u64,
22        ) -> Pin<Box<dyn Future<Output = Result<ItemCollection>> + Send>>
23        + Send,
24>;
25
26/// A backend's streamed search: the item stream plus the finalizer for the collection footer.
27pub struct StreamedSearch {
28    /// The response items, streamed one at a time.
29    pub items: ItemStream,
30    /// Produces the finished collection once the items drain.
31    pub finalize: Finalize,
32}
33
34/// A backend that streams a search response as items plus a finished [`ItemCollection`]. How it produces
35/// and paginates them is the implementation's concern.
36///
37/// `context` requests `numberMatched` (the STAC context extension). `self_href` is the URL this response
38/// is served at; pagination links are absolute against it, or relative when `None`.
39pub trait StreamSearch: Send + Sync {
40    /// Begins a streamed search, capped at `max_items` total items.
41    fn stream_search(
42        &self,
43        search: Search,
44        max_items: Option<usize>,
45        context: bool,
46        self_href: Option<String>,
47    ) -> impl Future<Output = Result<StreamedSearch>> + Send;
48
49    /// Drives this backend's streamed search into `writer` as one flat-memory JSON `ItemCollection`,
50    /// returning the number of items written.
51    fn write_search<W: Write>(
52        &self,
53        search: Search,
54        max_items: Option<usize>,
55        context: bool,
56        self_href: Option<String>,
57        writer: W,
58        pretty: bool,
59    ) -> impl Future<Output = Result<u64>> {
60        async move {
61            let StreamedSearch { items, finalize } = self
62                .stream_search(search, max_items, context, self_href)
63                .await?;
64            write_item_collection(writer, items, pretty, finalize).await
65        }
66    }
67}
68
69/// Writes a search response as a streamed `FeatureCollection`: the `features` array is written one item
70/// at a time, then `finalize` supplies the footer (links + optional count). The bytes equal
71/// `serde_json`-serializing the equivalent [`ItemCollection`], pretty or compact. Returns the item count.
72pub async fn write_item_collection<W, S, F, Fut>(
73    mut writer: W,
74    items: S,
75    pretty: bool,
76    finalize: F,
77) -> Result<u64>
78where
79    W: Write,
80    S: Stream<Item = Result<Value>>,
81    F: FnOnce(Option<Value>, Option<Value>, u64) -> Fut,
82    Fut: Future<Output = Result<ItemCollection>>,
83{
84    writer.write_all(if pretty {
85        b"{\n  \"type\": \"FeatureCollection\",\n  \"features\": ["
86    } else {
87        b"{\"type\":\"FeatureCollection\",\"features\":["
88    })?;
89
90    futures::pin_mut!(items);
91    let mut first: Option<Value> = None;
92    let mut pending: Option<Value> = None;
93    let mut count: u64 = 0;
94    while let Some(item) = items.next().await {
95        let item = item?;
96        if let Some(previous) = pending.take() {
97            write_element(&mut writer, &previous, count, pretty)?;
98            count += 1;
99        } else {
100            first = Some(item.clone());
101        }
102        pending = Some(item);
103    }
104    if let Some(last) = &pending {
105        write_element(&mut writer, last, count, pretty)?;
106        count += 1;
107    }
108    writer.write_all(if pretty && count > 0 { b"\n  ]" } else { b"]" })?;
109
110    // The footer is the rest of a real ItemCollection (`links`, `numberMatched`, `numberReturned`, …),
111    // serialized by serde and spliced in after the streamed features — `type`/`features` dropped since
112    // they're already written.
113    let mut collection = finalize(first, pending, count).await?;
114    collection.number_returned = Some(count);
115    let value = serde_json::to_value(&collection)?;
116    let members: serde_json::Map<String, Value> = value
117        .as_object()
118        .expect("an ItemCollection serializes to a JSON object")
119        .iter()
120        .filter(|(key, _)| key.as_str() != "type" && key.as_str() != "features")
121        .map(|(key, value)| (key.clone(), value.clone()))
122        .collect();
123    if !members.is_empty() {
124        let object = if pretty {
125            serde_json::to_string_pretty(&Value::Object(members))?
126        } else {
127            serde_json::to_string(&Value::Object(members))?
128        };
129        let inner = object
130            .strip_prefix('{')
131            .and_then(|rest| rest.strip_suffix('}'))
132            .expect("serde_json serializes an object with braces");
133        writer.write_all(b",")?;
134        writer.write_all(inner.trim_end().as_bytes())?;
135    }
136
137    writer.write_all(if pretty { b"\n}" } else { b"}" })?;
138    Ok(count)
139}
140
141/// Writes one item as an element of the `features` array. `index` is the
142/// element's position (0-based); a non-zero index gets a leading separator.
143fn write_element<W: Write>(writer: &mut W, item: &Value, index: u64, pretty: bool) -> Result<()> {
144    if pretty {
145        writer.write_all(if index == 0 { b"\n" } else { b",\n" })?;
146        let element = serde_json::to_string_pretty(item)?;
147        for (line_index, line) in element.lines().enumerate() {
148            if line_index > 0 {
149                writer.write_all(b"\n")?;
150            }
151            // Elements sit two levels deep (indent 4) inside the root object.
152            writer.write_all(b"    ")?;
153            writer.write_all(line.as_bytes())?;
154        }
155    } else {
156        if index > 0 {
157            writer.write_all(b",")?;
158        }
159        serde_json::to_writer(&mut *writer, item)?;
160    }
161    Ok(())
162}
163
164#[cfg(test)]
165mod tests {
166    use super::write_item_collection;
167    use futures::stream;
168    use serde_json::Value;
169    use stac::{Item, Link, api::ItemCollection};
170
171    /// `n` serialized STAC items and the same items as api items, so the stream
172    /// input and the expected collection agree byte-for-byte.
173    fn items(n: usize) -> (Vec<Value>, Vec<stac::api::Item>) {
174        let api: Vec<stac::api::Item> = (0..n)
175            .map(|i| Item::new(format!("item-{i}")).try_into().unwrap())
176            .collect();
177        let values = api
178            .iter()
179            .map(|i| serde_json::to_value(i).unwrap())
180            .collect();
181        (values, api)
182    }
183
184    async fn run(
185        values: Vec<Value>,
186        links: Vec<Link>,
187        matched: Option<u64>,
188        pretty: bool,
189    ) -> Vec<u8> {
190        let footer_links = links;
191        let mut buf = Vec::new();
192        write_item_collection(
193            &mut buf,
194            stream::iter(values.into_iter().map(Ok)),
195            pretty,
196            |_first, _last, _count| async move {
197                let mut collection = ItemCollection::new(Vec::<stac::api::Item>::new()).unwrap();
198                collection.links = footer_links;
199                collection.number_matched = matched;
200                Ok(collection)
201            },
202        )
203        .await
204        .unwrap();
205        buf
206    }
207
208    #[tokio::test]
209    async fn byte_identical_to_buffered() {
210        let links = vec![Link::new("http://example.com/next?token=abc", "next")];
211        for n in [0usize, 1, 2, 5] {
212            for pretty in [false, true] {
213                let matched = Some(n as u64 + 100);
214                let (values, api) = items(n);
215                let got = run(values, links.clone(), matched, pretty).await;
216
217                let mut want_ic = ItemCollection::new(api).unwrap();
218                want_ic.links = links.clone();
219                want_ic.number_matched = matched;
220                let want = if pretty {
221                    serde_json::to_vec_pretty(&want_ic).unwrap()
222                } else {
223                    serde_json::to_vec(&want_ic).unwrap()
224                };
225                assert_eq!(
226                    String::from_utf8(got).unwrap(),
227                    String::from_utf8(want).unwrap(),
228                    "n={n} pretty={pretty}"
229                );
230            }
231        }
232    }
233}
stac_io/stream.rs

stac_io/
stream.rs
