stac_server/backend/
memory.rs

1use crate::{Backend, DEFAULT_LIMIT, Error, Result};
2use serde_json::Map;
3use stac::{Collection, Item};
4use stac_api::{ItemCollection, Items, Search};
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 Backend for MemoryBackend {
37    fn has_item_search(&self) -> bool {
38        true
39    }
40
41    fn has_filter(&self) -> bool {
42        false
43    }
44
45    async fn collections(&self) -> Result<Vec<Collection>> {
46        let collections = self.collections.read().unwrap();
47        Ok(collections.values().cloned().collect())
48    }
49
50    async fn collection(&self, id: &str) -> Result<Option<Collection>> {
51        let collections = self.collections.read().unwrap();
52        Ok(collections.get(id).cloned())
53    }
54
55    async fn add_collection(&mut self, collection: Collection) -> Result<()> {
56        let mut collections = self.collections.write().unwrap();
57        let _ = collections.insert(collection.id.clone(), collection);
58        Ok(())
59    }
60
61    async fn add_item(&mut self, item: Item) -> Result<()> {
62        if let Some(collection_id) = item.collection.clone() {
63            if self.collection(&collection_id).await?.is_none() {
64                Err(Error::MemoryBackend(format!(
65                    "no collection with id='{collection_id}'",
66                )))
67            } else {
68                let mut items = self.items.write().unwrap();
69                items.entry(collection_id).or_default().push(item);
70                Ok(())
71            }
72        } else {
73            Err(Error::MemoryBackend(format!(
74                "collection not set on item: {}",
75                item.id
76            )))
77        }
78    }
79
80    async fn items(&self, collection_id: &str, items: Items) -> Result<Option<ItemCollection>> {
81        {
82            let collections = self.collections.read().unwrap();
83            if !collections.contains_key(collection_id) {
84                return Ok(None);
85            }
86        };
87        let search = items.search_collection(collection_id);
88        self.search(search).await.map(Some)
89    }
90
91    async fn item(&self, collection_id: &str, item_id: &str) -> Result<Option<Item>> {
92        let items = self.items.read().unwrap();
93        Ok(items
94            .get(collection_id)
95            .and_then(|items| items.iter().find(|item| item.id == item_id).cloned()))
96    }
97
98    async fn search(&self, mut search: Search) -> Result<ItemCollection> {
99        let items = self.items.read().unwrap();
100        if search.collections.is_empty() {
101            search.collections = items.keys().cloned().collect();
102        }
103        let mut item_references = Vec::new();
104        for collection in &search.collections {
105            if let Some(items) = items.get(collection) {
106                item_references.extend(
107                    items
108                        .iter()
109                        .filter(|item| search.matches(item).unwrap_or_default()),
110                );
111            }
112        }
113        let limit = search.limit.unwrap_or(DEFAULT_LIMIT).try_into()?;
114        let skip = search
115            .additional_fields
116            .get("skip")
117            .and_then(|skip| skip.as_str())
118            .and_then(|skip| skip.parse::<u64>().ok())
119            .unwrap_or_default()
120            .try_into()?;
121        let len = item_references.len();
122        let items = item_references
123            .into_iter()
124            .skip(skip)
125            .take(limit)
126            .map(|item| stac_api::Item::try_from(item.clone()).map_err(Error::from))
127            .collect::<Result<Vec<_>>>()?;
128        let mut item_collection = ItemCollection::new(items)?;
129        if len > item_collection.items.len() + skip {
130            let mut next = Map::new();
131            let _ = next.insert("skip".to_string(), (skip + limit).into());
132            item_collection.next = Some(next);
133        }
134        if skip > 0 {
135            let mut prev = Map::new();
136            let skip = skip.saturating_sub(limit);
137            let _ = prev.insert("skip".to_string(), skip.into());
138            item_collection.prev = Some(prev);
139        }
140        Ok(item_collection)
141    }
142}
143
144impl Default for MemoryBackend {
145    fn default() -> Self {
146        Self::new()
147    }
148}