stac_server/backend/
pgstac.rs

1use crate::{Backend, Error, Result};
2use bb8::Pool;
3use bb8_postgres::PostgresConnectionManager;
4use pgstac::Pgstac;
5use rustls::{ClientConfig, RootCertStore};
6use serde_json::Map;
7use stac::{Collection, Item};
8use stac_api::{ItemCollection, Items, Search};
9use tokio_postgres::{
10    tls::{MakeTlsConnect, TlsConnect},
11    Socket,
12};
13use tokio_postgres_rustls::MakeRustlsConnect;
14
15/// A backend for a [pgstac](https://github.com/stac-utils/pgstac) database.
16#[derive(Clone, Debug)]
17pub struct PgstacBackend<Tls>
18where
19    Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
20    <Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
21    <Tls as MakeTlsConnect<Socket>>::TlsConnect: Send,
22    <<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
23{
24    pool: Pool<PostgresConnectionManager<Tls>>,
25}
26
27impl PgstacBackend<MakeRustlsConnect> {
28    /// Creates a new PgstacBackend from a string-like configuration.
29    ///
30    /// This will use an unverified tls. To provide your own tls, use
31    /// [PgstacBackend::new_from_stringlike_and_tls].
32    ///
33    /// # Examples
34    ///
35    /// ```no_run
36    /// use stac_server::PgstacBackend;
37    /// # tokio_test::block_on(async {
38    /// let backend = PgstacBackend::new_from_stringlike("postgresql://username:password@localhost:5432/postgis").await.unwrap();
39    /// # })
40    /// ```
41    pub async fn new_from_stringlike(
42        params: impl ToString,
43    ) -> Result<PgstacBackend<MakeRustlsConnect>> {
44        let config = ClientConfig::builder()
45            .with_root_certificates(RootCertStore::empty())
46            .with_no_client_auth();
47        let tls = MakeRustlsConnect::new(config);
48        PgstacBackend::new_from_stringlike_and_tls(params, tls).await
49    }
50}
51
52impl<Tls> PgstacBackend<Tls>
53where
54    Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
55    <Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
56    <Tls as MakeTlsConnect<Socket>>::TlsConnect: Send,
57    <<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
58{
59    /// Creates a new PgstacBackend from a string-like configuration and a tls.
60    pub async fn new_from_stringlike_and_tls(
61        params: impl ToString,
62        tls: Tls,
63    ) -> Result<PgstacBackend<Tls>> {
64        let params = params.to_string();
65        let connection_manager = PostgresConnectionManager::new_from_stringlike(params, tls)?;
66        let pool = Pool::builder().build(connection_manager).await?;
67        Ok(PgstacBackend { pool })
68    }
69}
70
71impl<Tls> Backend for PgstacBackend<Tls>
72where
73    Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
74    <Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
75    <Tls as MakeTlsConnect<Socket>>::TlsConnect: Send,
76    <<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
77{
78    fn has_item_search(&self) -> bool {
79        true
80    }
81
82    fn has_filter(&self) -> bool {
83        true
84    }
85
86    async fn add_collection(&mut self, collection: Collection) -> Result<()> {
87        let client = self.pool.get().await?;
88        client.add_collection(collection).await.map_err(Error::from)
89    }
90
91    async fn collection(&self, id: &str) -> Result<Option<Collection>> {
92        let client = self.pool.get().await?;
93        let value = client.collection(id).await?;
94        value
95            .map(serde_json::from_value)
96            .transpose()
97            .map_err(Error::from)
98    }
99
100    async fn collections(&self) -> Result<Vec<Collection>> {
101        let client = self.pool.get().await?;
102        let values = client.collections().await?;
103        values
104            .into_iter()
105            .map(|v| serde_json::from_value(v).map_err(Error::from))
106            .collect()
107    }
108
109    async fn add_item(&mut self, item: Item) -> Result<()> {
110        let client = self.pool.get().await?;
111        client.add_item(item).await.map_err(Error::from)
112    }
113
114    async fn add_items(&mut self, items: Vec<Item>) -> Result<()> {
115        tracing::debug!("adding {} items using pgstac loading", items.len());
116        let client = self.pool.get().await?;
117        client.add_items(&items).await.map_err(Error::from)
118    }
119
120    async fn items(&self, collection_id: &str, items: Items) -> Result<Option<ItemCollection>> {
121        // TODO should we check for collection existence?
122        let search = items.search_collection(collection_id);
123        self.search(search).await.map(Some)
124    }
125
126    async fn item(&self, collection_id: &str, item_id: &str) -> Result<Option<Item>> {
127        let client = self.pool.get().await?;
128        let value = client.item(item_id, Some(collection_id)).await?;
129        value
130            .map(serde_json::from_value)
131            .transpose()
132            .map_err(Error::from)
133    }
134
135    async fn search(&self, search: Search) -> Result<ItemCollection> {
136        let client = self.pool.get().await?;
137        let page = client.search(search).await?;
138        let next_token = page.next_token();
139        let prev_token = page.prev_token();
140        let mut item_collection = ItemCollection::new(page.features)?;
141        if let Some(next_token) = next_token {
142            let mut next = Map::new();
143            let _ = next.insert("token".into(), next_token.into());
144            item_collection.next = Some(next);
145        }
146        if let Some(prev_token) = prev_token {
147            let mut prev = Map::new();
148            let _ = prev.insert("token".into(), prev_token.into());
149            item_collection.prev = Some(prev);
150        }
151        item_collection.context = page.context;
152        Ok(item_collection)
153    }
154}