bring streaming search traits and get ready for pgstac 0.10

Cargo.toml

@@ -74,7 +74,10 @@ mime = "0.3.17"
 mockito = "1.5"
 object_store = "0.13.0"
 parquet = { version = "58.0.0" }
-pgstac = "0.4.9"
+# pgstac 0.10, from the main pgstac tree (v010io). One crate serves both roles: streaming
+# `search --pgstac` (the `search-writer` feature) and the pgstac `serve` backend (the `pool` feature) —
+# consumers pick. Shares this workspace's stac + stac-io (see [patch.crates-io]).
+pgstac = { version = "0.9.11-dev", path = "/home/bitner/data/pgstac/src/pgstac-rs", default-features = false }
 quote = "1.0"
 referencing = { version = "0.47.0", features = ["retrieve-async"] }
 reqwest = { version = "0.13.1", features = ["query"] }
@@ -105,3 +108,6 @@ wkb = "0.9.0"
 
 [patch.crates-io]
 stac = { path = "crates/core" }
+# So pgstac 0.10 shares this workspace's stac-io — the one with the streaming ItemCollection writer —
+# instead of resolving a separate published copy.
+stac-io = { path = "crates/io" }

crates/cli/Cargo.toml

@@ -13,6 +13,8 @@ rust-version.workspace = true
 
 [features]
 default = []
+# pgstac support: streaming `search --pgstac` (the search-writer StreamSearch backend) + pgstac-backed
+# `serve` (via stac-server). Both resolve to the one pgstac 0.10 crate.
 pgstac = ["dep:pgstac", "stac-server/pgstac"]
 duckdb-bundled = ["stac-duckdb/bundled"]
 
@@ -24,7 +26,7 @@ clap = { workspace = true, features = ["derive"] }
 clap_complete.workspace = true
 futures-core.workspace = true
 futures-util.workspace = true
-pgstac = { workspace = true, optional = true }
+pgstac = { workspace = true, optional = true, features = ["search-writer"] }
 serde_json.workspace = true
 stac = { version = "0.17.2", path = "../core" }
 stac-duckdb = { version = "0.3.9", path = "../duckdb" }

crates/cli/src/lib.rs

@@ -247,6 +247,18 @@ pub enum Command {
             value_parser = |s: &str| KeyValue::from_str(s).map(|kv| (kv.0, kv.1))
         )]
         headers: Option<HeaderMap>,
+
+        /// Request the total match count (the `context` extension's `numberMatched`).
+        ///
+        /// Only the `postgresql` implementation honors this; it costs a second query.
+        #[arg(long = "context")]
+        context: bool,
+
+        /// Base URL that pagination (`next`/`prev`) links extend, i.e. the search endpoint's own
+        /// (`self`) URL. Without it, links are relative (`?token=…`). Only the `postgresql`
+        /// implementation emits pagination links.
+        #[arg(long = "self-href")]
+        self_href: Option<String>,
     },
 
     /// Serves a STAC API.
@@ -426,6 +438,8 @@ impl Rustac {
                 ref filter,
                 ref limit,
                 ref headers,
+                context,
+                ref self_href,
             } => {
                 // Infer the search implementation from the href if not explicitly provided
                 let search_impl = search_with.unwrap_or_else(|| {
@@ -459,11 +473,51 @@ impl Rustac {
                     SearchImplementation::Postgresql => {
                         #[cfg(feature = "pgstac")]
                         {
-                            pgstac::search(href, search, *max_items).await?
+                            use stac_io::StreamSearch as _;
+                            use std::io::Write as _;
+                            let backend = pgstac::PgstacPool::connect(pgstac::ConnectConfig {
+                                dsn: Some(href.to_string()),
+                                ..Default::default()
+                            })
+                            .await?;
+                            let dest = outfile
+                                .as_deref()
+                                .and_then(|path| if path == "-" { None } else { Some(path) });
+                            let pretty = matches!(self.output_format(dest), Format::Json(true));
+                            if let Some(path) = dest {
+                                let file = std::fs::File::create(path)?;
+                                backend
+                                    .write_search(
+                                        search,
+                                        *max_items,
+                                        context,
+                                        self_href.clone(),
+                                        file,
+                                        pretty,
+                                    )
+                                    .await?;
+                            } else {
+                                let stdout = std::io::stdout();
+                                let mut handle = stdout.lock();
+                                backend
+                                    .write_search(
+                                        search,
+                                        *max_items,
+                                        context,
+                                        self_href.clone(),
+                                        &mut handle,
+                                        pretty,
+                                    )
+                                    .await?;
+                                handle.flush()?;
+                            }
+                            return Ok(());
                         }
                         #[cfg(not(feature = "pgstac"))]
                         {
-                            return Err(anyhow!("rustac is not compiled with pgstac support"));
+                            return Err(anyhow!(
+                                "rustac is not compiled with pgstac search support (enable the `pgstac` feature)"
+                            ));
                         }
                     }
                     SearchImplementation::Duckdb => stac_duckdb::search(href, search, *max_items)?,

crates/io/src/error.rs

@@ -4,6 +4,10 @@ use thiserror::Error;
 #[derive(Error, Debug)]
 #[non_exhaustive]
 pub enum Error {
+    /// An error from a [streaming search](crate::stream) backend, which produces its own error type.
+    #[error(transparent)]
+    Backend(#[from] Box<dyn std::error::Error + Send + Sync>),
+
     /// Returned when unable to read a STAC value from a path.
     #[error("{io}: {path}")]
     FromPath {

crates/io/src/lib.rs

@@ -9,6 +9,7 @@ mod read;
 mod realized_href;
 #[cfg(feature = "store")]
 pub mod store;
+pub mod stream;
 mod write;
 
 #[cfg(feature = "geoparquet")]
@@ -22,6 +23,7 @@ pub use {
     ndjson::{FromNdjsonPath, ToNdjsonPath, ndjson_item_reader},
     read::read,
     realized_href::RealizedHref,
+    stream::{Finalize, ItemStream, StreamSearch, StreamedSearch, write_item_collection},
     write::write,
 };
 

crates/io/src/stream.rs

@@ -0,0 +1,233 @@
+//! Streaming JSON writer for a search-response [`ItemCollection`](stac::api::ItemCollection): the
+//! `features` array is written one item at a time, and the rest of the collection (links, context,
+//! counts) is supplied by a `finalize` callback after the items drain (the `next` link needs the last
+//! item; a `numberMatched` count may run concurrently).
+
+use crate::Result;
+use futures::{Stream, StreamExt};
+use serde_json::Value;
+use stac::api::{ItemCollection, Search};
+use std::{future::Future, io::Write, pin::Pin};
+
+/// A boxed, pinned stream of serialized STAC items.
+pub type ItemStream = Pin<Box<dyn Stream<Item = Result<Value>> + Send>>;
+
+/// Produces the finished [`ItemCollection`] (with empty `items`; the writer fills `numberReturned`) from
+/// the first item, the last item, and the number written. Called once, after the stream drains.
+pub type Finalize = Box<
+    dyn FnOnce(
+            Option<Value>,
+            Option<Value>,
+            u64,
+        ) -> Pin<Box<dyn Future<Output = Result<ItemCollection>> + Send>>
+        + Send,
+>;
+
+/// A backend's streamed search: the item stream plus the finalizer for the collection footer.
+pub struct StreamedSearch {
+    /// The response items, streamed one at a time.
+    pub items: ItemStream,
+    /// Produces the finished collection once the items drain.
+    pub finalize: Finalize,
+}
+
+/// A backend that streams a search response as items plus a finished [`ItemCollection`]. How it produces
+/// and paginates them is the implementation's concern.
+///
+/// `context` requests `numberMatched` (the STAC context extension). `self_href` is the URL this response
+/// is served at; pagination links are absolute against it, or relative when `None`.
+pub trait StreamSearch: Send + Sync {
+    /// Begins a streamed search, capped at `max_items` total items.
+    fn stream_search(
+        &self,
+        search: Search,
+        max_items: Option<usize>,
+        context: bool,
+        self_href: Option<String>,
+    ) -> impl Future<Output = Result<StreamedSearch>> + Send;
+
+    /// Drives this backend's streamed search into `writer` as one flat-memory JSON `ItemCollection`,
+    /// returning the number of items written.
+    fn write_search<W: Write>(
+        &self,
+        search: Search,
+        max_items: Option<usize>,
+        context: bool,
+        self_href: Option<String>,
+        writer: W,
+        pretty: bool,
+    ) -> impl Future<Output = Result<u64>> {
+        async move {
+            let StreamedSearch { items, finalize } = self
+                .stream_search(search, max_items, context, self_href)
+                .await?;
+            write_item_collection(writer, items, pretty, finalize).await
+        }
+    }
+}
+
+/// Writes a search response as a streamed `FeatureCollection`: the `features` array is written one item
+/// at a time, then `finalize` supplies the footer (links + optional count). The bytes equal
+/// `serde_json`-serializing the equivalent [`ItemCollection`], pretty or compact. Returns the item count.
+pub async fn write_item_collection<W, S, F, Fut>(
+    mut writer: W,
+    items: S,
+    pretty: bool,
+    finalize: F,
+) -> Result<u64>
+where
+    W: Write,
+    S: Stream<Item = Result<Value>>,
+    F: FnOnce(Option<Value>, Option<Value>, u64) -> Fut,
+    Fut: Future<Output = Result<ItemCollection>>,
+{
+    writer.write_all(if pretty {
+        b"{\n  \"type\": \"FeatureCollection\",\n  \"features\": ["
+    } else {
+        b"{\"type\":\"FeatureCollection\",\"features\":["
+    })?;
+
+    futures::pin_mut!(items);
+    let mut first: Option<Value> = None;
+    let mut pending: Option<Value> = None;
+    let mut count: u64 = 0;
+    while let Some(item) = items.next().await {
+        let item = item?;
+        if let Some(previous) = pending.take() {
+            write_element(&mut writer, &previous, count, pretty)?;
+            count += 1;
+        } else {
+            first = Some(item.clone());
+        }
+        pending = Some(item);
+    }
+    if let Some(last) = &pending {
+        write_element(&mut writer, last, count, pretty)?;
+        count += 1;
+    }
+    writer.write_all(if pretty && count > 0 { b"\n  ]" } else { b"]" })?;
+
+    // The footer is the rest of a real ItemCollection (`links`, `numberMatched`, `numberReturned`, …),
+    // serialized by serde and spliced in after the streamed features — `type`/`features` dropped since
+    // they're already written.
+    let mut collection = finalize(first, pending, count).await?;
+    collection.number_returned = Some(count);
+    let value = serde_json::to_value(&collection)?;
+    let members: serde_json::Map<String, Value> = value
+        .as_object()
+        .expect("an ItemCollection serializes to a JSON object")
+        .iter()
+        .filter(|(key, _)| key.as_str() != "type" && key.as_str() != "features")
+        .map(|(key, value)| (key.clone(), value.clone()))
+        .collect();
+    if !members.is_empty() {
+        let object = if pretty {
+            serde_json::to_string_pretty(&Value::Object(members))?
+        } else {
+            serde_json::to_string(&Value::Object(members))?
+        };
+        let inner = object
+            .strip_prefix('{')
+            .and_then(|rest| rest.strip_suffix('}'))
+            .expect("serde_json serializes an object with braces");
+        writer.write_all(b",")?;
+        writer.write_all(inner.trim_end().as_bytes())?;
+    }
+
+    writer.write_all(if pretty { b"\n}" } else { b"}" })?;
+    Ok(count)
+}
+
+/// Writes one item as an element of the `features` array. `index` is the
+/// element's position (0-based); a non-zero index gets a leading separator.
+fn write_element<W: Write>(writer: &mut W, item: &Value, index: u64, pretty: bool) -> Result<()> {
+    if pretty {
+        writer.write_all(if index == 0 { b"\n" } else { b",\n" })?;
+        let element = serde_json::to_string_pretty(item)?;
+        for (line_index, line) in element.lines().enumerate() {
+            if line_index > 0 {
+                writer.write_all(b"\n")?;
+            }
+            // Elements sit two levels deep (indent 4) inside the root object.
+            writer.write_all(b"    ")?;
+            writer.write_all(line.as_bytes())?;
+        }
+    } else {
+        if index > 0 {
+            writer.write_all(b",")?;
+        }
+        serde_json::to_writer(&mut *writer, item)?;
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::write_item_collection;
+    use futures::stream;
+    use serde_json::Value;
+    use stac::{Item, Link, api::ItemCollection};
+
+    /// `n` serialized STAC items and the same items as api items, so the stream
+    /// input and the expected collection agree byte-for-byte.
+    fn items(n: usize) -> (Vec<Value>, Vec<stac::api::Item>) {
+        let api: Vec<stac::api::Item> = (0..n)
+            .map(|i| Item::new(format!("item-{i}")).try_into().unwrap())
+            .collect();
+        let values = api
+            .iter()
+            .map(|i| serde_json::to_value(i).unwrap())
+            .collect();
+        (values, api)
+    }
+
+    async fn run(
+        values: Vec<Value>,
+        links: Vec<Link>,
+        matched: Option<u64>,
+        pretty: bool,
+    ) -> Vec<u8> {
+        let footer_links = links;
+        let mut buf = Vec::new();
+        write_item_collection(
+            &mut buf,
+            stream::iter(values.into_iter().map(Ok)),
+            pretty,
+            |_first, _last, _count| async move {
+                let mut collection = ItemCollection::new(Vec::<stac::api::Item>::new()).unwrap();
+                collection.links = footer_links;
+                collection.number_matched = matched;
+                Ok(collection)
+            },
+        )
+        .await
+        .unwrap();
+        buf
+    }
+
+    #[tokio::test]
+    async fn byte_identical_to_buffered() {
+        let links = vec![Link::new("http://example.com/next?token=abc", "next")];
+        for n in [0usize, 1, 2, 5] {
+            for pretty in [false, true] {
+                let matched = Some(n as u64 + 100);
+                let (values, api) = items(n);
+                let got = run(values, links.clone(), matched, pretty).await;
+
+                let mut want_ic = ItemCollection::new(api).unwrap();
+                want_ic.links = links.clone();
+                want_ic.number_matched = matched;
+                let want = if pretty {
+                    serde_json::to_vec_pretty(&want_ic).unwrap()
+                } else {
+                    serde_json::to_vec(&want_ic).unwrap()
+                };
+                assert_eq!(
+                    String::from_utf8(got).unwrap(),
+                    String::from_utf8(want).unwrap(),
+                    "n={n} pretty={pretty}"
+                );
+            }
+        }
+    }
+}