Skip to main content

stac_server/backend/
memory.rs

1use crate::{Backend, DEFAULT_LIMIT, Error, Result};
2use serde_json::Map;
3use stac::api::{CollectionSearchClient, ItemCollection, Search, SearchClient, TransactionClient};
4use stac::{Collection, Item};
5use std::{
6    collections::{BTreeMap, HashMap},
7    sync::{Arc, RwLock},
8};
9
10/// A naive backend that stores collections and items in memory.
11///
12/// This backend is meant to be used for testing and toy servers, not for production.
13#[derive(Clone, Debug)]
14pub struct MemoryBackend {
15    collections: Arc<RwLock<BTreeMap<String, Collection>>>,
16    items: Arc<RwLock<HashMap<String, Vec<Item>>>>,
17}
18
19impl MemoryBackend {
20    /// Creates a new memory backend.
21    ///
22    /// # Examples
23    ///
24    /// ```
25    /// use stac_server::MemoryBackend;
26    /// let backend = MemoryBackend::new();
27    /// ```
28    pub fn new() -> MemoryBackend {
29        MemoryBackend {
30            collections: Arc::new(RwLock::new(BTreeMap::new())),
31            items: Arc::new(RwLock::new(HashMap::new())),
32        }
33    }
34}
35
36impl SearchClient for MemoryBackend {
37    type Error = Error;
38
39    async fn search(&self, mut search: Search) -> Result<ItemCollection> {
40        let items = self.items.read().unwrap();
41        if search.collections.is_empty() {
42            search.collections = items.keys().cloned().collect();
43        }
44        let mut item_references = Vec::new();
45        for collection in &search.collections {
46            if let Some(items) = items.get(collection) {
47                item_references.extend(
48                    items
49                        .iter()
50                        .filter(|item| search.matches(item).unwrap_or_default()),
51                );
52            }
53        }
54        let limit = search.limit.unwrap_or(DEFAULT_LIMIT).try_into()?;
55        let skip = search
56            .additional_fields
57            .get("skip")
58            .and_then(|skip| skip.as_str())
59            .and_then(|skip| skip.parse::<u64>().ok())
60            .unwrap_or_default()
61            .try_into()?;
62        let len = item_references.len();
63        let items = item_references
64            .into_iter()
65            .skip(skip)
66            .take(limit)
67            .map(|item| stac::api::Item::try_from(item.clone()).map_err(Error::from))
68            .collect::<Result<Vec<_>>>()?;
69        let mut item_collection = ItemCollection::new(items)?;
70        if len > item_collection.items.len() + skip {
71            let mut next = Map::new();
72            let _ = next.insert("skip".to_string(), (skip + limit).into());
73            item_collection.next = Some(next);
74        }
75        if skip > 0 {
76            let mut prev = Map::new();
77            let skip = skip.saturating_sub(limit);
78            let _ = prev.insert("skip".to_string(), skip.into());
79            item_collection.prev = Some(prev);
80        }
81        Ok(item_collection)
82    }
83
84    async fn item(&self, collection_id: &str, item_id: &str) -> Result<Option<Item>> {
85        let items = self.items.read().unwrap();
86        Ok(items
87            .get(collection_id)
88            .and_then(|items| items.iter().find(|item| item.id == item_id).cloned()))
89    }
90}
91
92impl CollectionSearchClient for MemoryBackend {
93    type Error = Error;
94
95    async fn collections(&self) -> Result<Vec<Collection>> {
96        let collections = self.collections.read().unwrap();
97        Ok(collections.values().cloned().collect())
98    }
99
100    async fn collection(&self, id: &str) -> Result<Option<Collection>> {
101        let collections = self.collections.read().unwrap();
102        Ok(collections.get(id).cloned())
103    }
104}
105
106impl TransactionClient for MemoryBackend {
107    type Error = Error;
108
109    async fn add_collection(&mut self, collection: Collection) -> Result<()> {
110        let mut collections = self.collections.write().unwrap();
111        let _ = collections.insert(collection.id.clone(), collection);
112        Ok(())
113    }
114
115    async fn add_item(&mut self, item: Item) -> Result<()> {
116        if let Some(collection_id) = item.collection.clone() {
117            if CollectionSearchClient::collection(self, &collection_id)
118                .await?
119                .is_none()
120            {
121                Err(Error::MemoryBackend(format!(
122                    "no collection with id='{collection_id}'",
123                )))
124            } else {
125                let mut items = self.items.write().unwrap();
126                items.entry(collection_id).or_default().push(item);
127                Ok(())
128            }
129        } else {
130            Err(Error::MemoryBackend(format!(
131                "collection not set on item: {}",
132                item.id
133            )))
134        }
135    }
136}
137
138impl Backend for MemoryBackend {
139    fn has_item_search(&self) -> bool {
140        true
141    }
142
143    fn has_filter(&self) -> bool {
144        false
145    }
146}
147
148impl Default for MemoryBackend {
149    fn default() -> Self {
150        Self::new()
151    }
152}