Skip to main content

stac_server/backend/
duckdb.rs

1use super::Backend;
2use crate::{Error, Result};
3use bb8::{ManageConnection, Pool};
4use stac::Collection;
5use stac::api::{CollectionSearchClient, Search, SearchClient, TransactionClient};
6use stac_duckdb::Client;
7
8/// A backend that uses [DuckDB](https://duckdb.org/) to query
9/// [stac-geoparquet](https://github.com/stac-utils/stac-geoparquet).
10#[derive(Clone, Debug)]
11pub struct DuckdbBackend {
12    pool: Pool<DuckdbConnectionManager>,
13}
14
15struct DuckdbConnectionManager {
16    href: String,
17}
18
19struct DuckdbConnection {
20    client: Client,
21    href: String,
22}
23
24impl DuckdbBackend {
25    /// Creates a new DuckDB backend pointing to a single **stac-geoparquet** file.
26    ///
27    /// # Examples
28    ///
29    /// ```
30    /// use stac_server::DuckdbBackend;
31    /// # tokio_test::block_on(async {
32    /// let backend = DuckdbBackend::new("data/100-sentinel-2-items.parquet").await.unwrap();
33    /// # })
34    /// ```
35    pub async fn new(href: impl ToString) -> Result<DuckdbBackend> {
36        let pool = Pool::builder()
37            .build(DuckdbConnectionManager {
38                href: href.to_string(),
39            })
40            .await?;
41        Ok(DuckdbBackend { pool })
42    }
43}
44
45impl SearchClient for DuckdbBackend {
46    type Error = Error;
47
48    async fn search(&self, search: Search) -> Result<stac::api::ItemCollection> {
49        let client = self.pool.get().await.map_err(Box::new)?;
50        client.search(search)
51    }
52}
53
54impl CollectionSearchClient for DuckdbBackend {
55    type Error = Error;
56
57    async fn collections(&self) -> Result<Vec<Collection>> {
58        let client = self.pool.get().await.map_err(Box::new)?;
59        client.collections()
60    }
61
62    async fn collection(&self, id: &str) -> Result<Option<Collection>> {
63        let client = self.pool.get().await.map_err(Box::new)?;
64        client.collection(id)
65    }
66}
67
68impl TransactionClient for DuckdbBackend {
69    type Error = Error;
70
71    async fn add_collection(&mut self, _collection: Collection) -> Result<()> {
72        Err(Error::ReadOnly)
73    }
74
75    async fn add_item(&mut self, _item: stac::Item) -> Result<()> {
76        Err(Error::ReadOnly)
77    }
78}
79
80impl Backend for DuckdbBackend {
81    fn has_item_search(&self) -> bool {
82        true
83    }
84
85    fn has_filter(&self) -> bool {
86        false
87    }
88}
89
90impl ManageConnection for DuckdbConnectionManager {
91    type Connection = DuckdbConnection;
92    type Error = Error;
93
94    async fn connect(&self) -> Result<DuckdbConnection> {
95        DuckdbConnection::new(&self.href)
96    }
97
98    async fn is_valid(&self, _conn: &mut DuckdbConnection) -> Result<()> {
99        Ok(())
100    }
101
102    fn has_broken(&self, _conn: &mut DuckdbConnection) -> bool {
103        false
104    }
105}
106
107impl DuckdbConnection {
108    fn new(href: impl ToString) -> Result<DuckdbConnection> {
109        let client = Client::new()?;
110        Ok(DuckdbConnection {
111            client,
112            href: href.to_string(),
113        })
114    }
115
116    fn collections(&self) -> Result<Vec<Collection>> {
117        let collections = self.client.collections(&self.href)?;
118        Ok(collections)
119    }
120
121    fn collection(&self, id: &str) -> Result<Option<Collection>> {
122        let collections = self.client.collections(&self.href)?;
123        Ok(collections
124            .into_iter()
125            .find(|collection| collection.id == id))
126    }
127
128    fn search(&self, search: Search) -> Result<stac::api::ItemCollection> {
129        let item_collection = self.client.search(&self.href, search)?;
130        Ok(item_collection)
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use stac::api::CollectionSearchClient;
137
138    #[tokio::test]
139    async fn backend() {
140        let backend = super::DuckdbBackend::new("data/100-sentinel-2-items.parquet")
141            .await
142            .unwrap();
143        assert!(
144            backend
145                .collection("sentinel-2-l2a")
146                .await
147                .unwrap()
148                .is_some()
149        );
150    }
151}