1use crate::storage::{StorageError, StorageKey, StoragePrefix, StorageProvider};
60use serde_json::Value;
61use std::collections::HashMap;
62use std::sync::Arc;
63use tokio::sync::RwLock;
64
65#[derive(Clone)]
72pub struct InMemoryStorage {
73 data: Arc<RwLock<HashMap<String, HashMap<String, HashMap<String, Value>>>>>,
75}
76
77impl InMemoryStorage {
78 pub fn new() -> Self {
80 Self {
81 data: Arc::new(RwLock::new(HashMap::new())),
82 }
83 }
84
85 pub async fn stats(&self) -> InMemoryStorageStats {
87 let data_guard = self.data.read().await;
88 let mut tenant_count = 0;
89 let mut resource_type_count = 0;
90 let mut total_resources = 0;
91
92 for (_, tenant_data) in data_guard.iter() {
93 tenant_count += 1;
94 for (_, type_data) in tenant_data.iter() {
95 resource_type_count += 1;
96 total_resources += type_data.len();
97 }
98 }
99
100 InMemoryStorageStats {
101 tenant_count,
102 resource_type_count,
103 total_resources,
104 }
105 }
106
107
108
109 fn extract_attribute_value(data: &Value, attribute_path: &str) -> Option<String> {
111 let parts: Vec<&str> = attribute_path.split('.').collect();
112 let mut current = data;
113
114 for part in parts {
115 if let Ok(index) = part.parse::<usize>() {
116 current = current.get(index)?;
118 } else {
119 current = current.get(part)?;
121 }
122 }
123
124 match current {
126 Value::String(s) => Some(s.clone()),
127 Value::Number(n) => Some(n.to_string()),
128 Value::Bool(b) => Some(b.to_string()),
129 _ => current.as_str().map(|s| s.to_string()),
130 }
131 }
132}
133
134impl Default for InMemoryStorage {
135 fn default() -> Self {
136 Self::new()
137 }
138}
139
140impl StorageProvider for InMemoryStorage {
141 type Error = StorageError;
142
143 async fn put(&self, key: StorageKey, data: Value) -> Result<Value, Self::Error> {
144 let mut data_guard = self.data.write().await;
145
146 let tenant_data = data_guard
148 .entry(key.tenant_id().to_string())
149 .or_insert_with(HashMap::new);
150
151 let type_data = tenant_data
152 .entry(key.resource_type().to_string())
153 .or_insert_with(HashMap::new);
154
155 type_data.insert(key.resource_id().to_string(), data.clone());
157
158 Ok(data)
160 }
161
162 async fn get(&self, key: StorageKey) -> Result<Option<Value>, Self::Error> {
163 let data_guard = self.data.read().await;
164
165 let result = data_guard
166 .get(key.tenant_id())
167 .and_then(|tenant_data| tenant_data.get(key.resource_type()))
168 .and_then(|type_data| type_data.get(key.resource_id()))
169 .cloned();
170
171 Ok(result)
172 }
173
174 async fn delete(&self, key: StorageKey) -> Result<bool, Self::Error> {
175 let mut data_guard = self.data.write().await;
176
177 let existed = if let Some(tenant_data) = data_guard.get_mut(key.tenant_id()) {
178 if let Some(type_data) = tenant_data.get_mut(key.resource_type()) {
179 type_data.remove(key.resource_id()).is_some()
180 } else {
181 false
182 }
183 } else {
184 false
185 };
186
187 Ok(existed)
188 }
189
190 async fn list(
191 &self,
192 prefix: StoragePrefix,
193 offset: usize,
194 limit: usize,
195 ) -> Result<Vec<(StorageKey, Value)>, Self::Error> {
196 if limit == 0 {
197 return Ok(Vec::new());
198 }
199
200 let data_guard = self.data.read().await;
201
202 let type_data = match data_guard
203 .get(prefix.tenant_id())
204 .and_then(|tenant_data| tenant_data.get(prefix.resource_type()))
205 {
206 Some(data) => data,
207 None => return Ok(Vec::new()),
208 };
209
210 let mut keys: Vec<_> = type_data.keys().collect();
212 keys.sort();
213
214 let results: Vec<(StorageKey, Value)> = keys
216 .into_iter()
217 .skip(offset)
218 .take(limit)
219 .filter_map(|resource_id| {
220 type_data.get(resource_id).map(|data| {
221 (
222 StorageKey::new(prefix.tenant_id(), prefix.resource_type(), resource_id),
223 data.clone(),
224 )
225 })
226 })
227 .collect();
228
229 Ok(results)
230 }
231
232 async fn find_by_attribute(
233 &self,
234 prefix: StoragePrefix,
235 attribute: &str,
236 value: &str,
237 ) -> Result<Vec<(StorageKey, Value)>, Self::Error> {
238 let data_guard = self.data.read().await;
239
240 let type_data = match data_guard
241 .get(prefix.tenant_id())
242 .and_then(|tenant_data| tenant_data.get(prefix.resource_type()))
243 {
244 Some(data) => data,
245 None => return Ok(Vec::new()),
246 };
247
248 let mut results = Vec::new();
249
250 for (resource_id, resource_data) in type_data {
251 if let Some(attr_value) = Self::extract_attribute_value(resource_data, attribute) {
252 if attr_value == value {
253 results.push((
254 StorageKey::new(prefix.tenant_id(), prefix.resource_type(), resource_id),
255 resource_data.clone(),
256 ));
257 }
258 }
259 }
260
261 results.sort_by(|a, b| a.0.resource_id().cmp(b.0.resource_id()));
263
264 Ok(results)
265 }
266
267 async fn exists(&self, key: StorageKey) -> Result<bool, Self::Error> {
268 let data_guard = self.data.read().await;
269
270 let exists = data_guard
271 .get(key.tenant_id())
272 .and_then(|tenant_data| tenant_data.get(key.resource_type()))
273 .and_then(|type_data| type_data.get(key.resource_id()))
274 .is_some();
275
276 Ok(exists)
277 }
278
279 async fn count(&self, prefix: StoragePrefix) -> Result<usize, Self::Error> {
280 let data_guard = self.data.read().await;
281
282 let count = data_guard
283 .get(prefix.tenant_id())
284 .and_then(|tenant_data| tenant_data.get(prefix.resource_type()))
285 .map(|type_data| type_data.len())
286 .unwrap_or(0);
287
288 Ok(count)
289 }
290
291 async fn list_tenants(&self) -> Result<Vec<String>, Self::Error> {
292 let data_guard = self.data.read().await;
293 Ok(data_guard.keys().cloned().collect())
294 }
295
296 async fn list_resource_types(&self, tenant_id: &str) -> Result<Vec<String>, Self::Error> {
297 let data_guard = self.data.read().await;
298 Ok(data_guard
299 .get(tenant_id)
300 .map(|tenant_data| tenant_data.keys().cloned().collect())
301 .unwrap_or_default())
302 }
303
304 async fn list_all_resource_types(&self) -> Result<Vec<String>, Self::Error> {
305 let data_guard = self.data.read().await;
306 let mut resource_types = std::collections::HashSet::new();
307
308 for tenant_data in data_guard.values() {
309 for resource_type in tenant_data.keys() {
310 resource_types.insert(resource_type.clone());
311 }
312 }
313
314 Ok(resource_types.into_iter().collect())
315 }
316
317 async fn clear(&self) -> Result<(), Self::Error> {
318 let mut data_guard = self.data.write().await;
319 data_guard.clear();
320 Ok(())
321 }
322}
323
324#[derive(Debug, Clone, PartialEq, Eq)]
326pub struct InMemoryStorageStats {
327 pub tenant_count: usize,
329 pub resource_type_count: usize,
331 pub total_resources: usize,
333}
334
335#[cfg(test)]
336mod tests {
337 use super::*;
338 use serde_json::json;
339
340 #[tokio::test]
341 async fn test_put_and_get() {
342 let storage = InMemoryStorage::new();
343 let key = StorageKey::new("tenant1", "User", "123");
344 let data = json!({"id": "123", "name": "test"});
345
346 let stored = storage.put(key.clone(), data.clone()).await.unwrap();
348 assert_eq!(stored, data);
349
350 let retrieved = storage.get(key).await.unwrap();
352 assert_eq!(retrieved, Some(data));
353 }
354
355 #[tokio::test]
356 async fn test_get_nonexistent() {
357 let storage = InMemoryStorage::new();
358 let key = StorageKey::new("tenant1", "User", "999");
359
360 let result = storage.get(key).await.unwrap();
361 assert!(result.is_none());
362 }
363
364 #[tokio::test]
365 async fn test_delete() {
366 let storage = InMemoryStorage::new();
367 let key = StorageKey::new("tenant1", "User", "123");
368 let data = json!({"id": "123", "name": "test"});
369
370 storage.put(key.clone(), data).await.unwrap();
372
373 let deleted = storage.delete(key.clone()).await.unwrap();
375 assert!(deleted);
376
377 let retrieved = storage.get(key.clone()).await.unwrap();
379 assert!(retrieved.is_none());
380
381 let deleted_again = storage.delete(key).await.unwrap();
383 assert!(!deleted_again);
384 }
385
386 #[tokio::test]
387 async fn test_exists() {
388 let storage = InMemoryStorage::new();
389 let key = StorageKey::new("tenant1", "User", "123");
390 let data = json!({"id": "123", "name": "test"});
391
392 assert!(!storage.exists(key.clone()).await.unwrap());
394
395 storage.put(key.clone(), data).await.unwrap();
397
398 assert!(storage.exists(key.clone()).await.unwrap());
400
401 storage.delete(key.clone()).await.unwrap();
403
404 assert!(!storage.exists(key).await.unwrap());
406 }
407
408 #[tokio::test]
409 async fn test_list_with_pagination() {
410 let storage = InMemoryStorage::new();
411 let prefix = StorageKey::prefix("tenant1", "User");
412
413 for i in 1..=5 {
415 let key = StorageKey::new("tenant1", "User", &format!("{}", i));
416 let data = json!({"id": i, "name": format!("user{}", i)});
417 storage.put(key, data).await.unwrap();
418 }
419
420 let page1 = storage.list(prefix.clone(), 0, 2).await.unwrap();
422 assert_eq!(page1.len(), 2);
423 assert_eq!(page1[0].0.resource_id(), "1");
424 assert_eq!(page1[1].0.resource_id(), "2");
425
426 let page2 = storage.list(prefix.clone(), 2, 2).await.unwrap();
427 assert_eq!(page2.len(), 2);
428 assert_eq!(page2[0].0.resource_id(), "3");
429 assert_eq!(page2[1].0.resource_id(), "4");
430
431 let page3 = storage.list(prefix, 4, 2).await.unwrap();
432 assert_eq!(page3.len(), 1);
433 assert_eq!(page3[0].0.resource_id(), "5");
434 }
435
436 #[tokio::test]
437 async fn test_find_by_attribute() {
438 let storage = InMemoryStorage::new();
439 let prefix = StorageKey::prefix("tenant1", "User");
440
441 let user1 = json!({
443 "id": "1",
444 "userName": "john.doe",
445 "emails": [{"value": "john@example.com", "primary": true}]
446 });
447 let user2 = json!({
448 "id": "2",
449 "userName": "jane.doe",
450 "emails": [{"value": "jane@example.com", "primary": true}]
451 });
452
453 storage
454 .put(StorageKey::new("tenant1", "User", "1"), user1)
455 .await
456 .unwrap();
457 storage
458 .put(StorageKey::new("tenant1", "User", "2"), user2)
459 .await
460 .unwrap();
461
462 let found = storage
464 .find_by_attribute(prefix.clone(), "userName", "john.doe")
465 .await
466 .unwrap();
467 assert_eq!(found.len(), 1);
468 assert_eq!(found[0].0.resource_id(), "1");
469
470 let found = storage
472 .find_by_attribute(prefix.clone(), "emails.0.value", "jane@example.com")
473 .await
474 .unwrap();
475 assert_eq!(found.len(), 1);
476 assert_eq!(found[0].0.resource_id(), "2");
477
478 let found = storage
480 .find_by_attribute(prefix, "userName", "nonexistent")
481 .await
482 .unwrap();
483 assert_eq!(found.len(), 0);
484 }
485
486 #[tokio::test]
487 async fn test_count() {
488 let storage = InMemoryStorage::new();
489 let prefix = StorageKey::prefix("tenant1", "User");
490
491 assert_eq!(storage.count(prefix.clone()).await.unwrap(), 0);
493
494 for i in 1..=3 {
496 let key = StorageKey::new("tenant1", "User", &format!("{}", i));
497 let data = json!({"id": i});
498 storage.put(key, data).await.unwrap();
499 }
500
501 assert_eq!(storage.count(prefix).await.unwrap(), 3);
502 }
503
504 #[tokio::test]
505 async fn test_tenant_isolation() {
506 let storage = InMemoryStorage::new();
507
508 let key1 = StorageKey::new("tenant1", "User", "123");
510 let key2 = StorageKey::new("tenant2", "User", "123");
511 let data1 = json!({"tenant": "1"});
512 let data2 = json!({"tenant": "2"});
513
514 storage.put(key1.clone(), data1.clone()).await.unwrap();
515 storage.put(key2.clone(), data2.clone()).await.unwrap();
516
517 assert_eq!(storage.get(key1).await.unwrap(), Some(data1));
519 assert_eq!(storage.get(key2).await.unwrap(), Some(data2));
520
521 assert_eq!(
523 storage
524 .count(StorageKey::prefix("tenant1", "User"))
525 .await
526 .unwrap(),
527 1
528 );
529 assert_eq!(
530 storage
531 .count(StorageKey::prefix("tenant2", "User"))
532 .await
533 .unwrap(),
534 1
535 );
536 }
537
538 #[tokio::test]
539 async fn test_stats() {
540 let storage = InMemoryStorage::new();
541
542 let stats = storage.stats().await;
544 assert_eq!(stats.tenant_count, 0);
545 assert_eq!(stats.resource_type_count, 0);
546 assert_eq!(stats.total_resources, 0);
547
548 storage
550 .put(StorageKey::new("tenant1", "User", "1"), json!({"id": "1"}))
551 .await
552 .unwrap();
553 storage
554 .put(StorageKey::new("tenant1", "User", "2"), json!({"id": "2"}))
555 .await
556 .unwrap();
557 storage
558 .put(StorageKey::new("tenant1", "Group", "1"), json!({"id": "1"}))
559 .await
560 .unwrap();
561 storage
562 .put(StorageKey::new("tenant2", "User", "1"), json!({"id": "1"}))
563 .await
564 .unwrap();
565
566 let stats = storage.stats().await;
567 assert_eq!(stats.tenant_count, 2);
568 assert_eq!(stats.resource_type_count, 3); assert_eq!(stats.total_resources, 4);
570 }
571
572 #[tokio::test]
573 async fn test_clear() {
574 let storage = InMemoryStorage::new();
575
576 storage
578 .put(StorageKey::new("tenant1", "User", "1"), json!({"id": "1"}))
579 .await
580 .unwrap();
581
582 assert_eq!(
584 storage
585 .count(StorageKey::prefix("tenant1", "User"))
586 .await
587 .unwrap(),
588 1
589 );
590
591 let _ = storage.clear().await;
593
594 assert_eq!(
596 storage
597 .count(StorageKey::prefix("tenant1", "User"))
598 .await
599 .unwrap(),
600 0
601 );
602 let stats = storage.stats().await;
603 assert_eq!(stats.total_resources, 0);
604 }
605
606 #[tokio::test]
607 async fn test_list_tenants_and_resource_types() {
608 let storage = InMemoryStorage::new();
609
610 storage
612 .put(StorageKey::new("tenant1", "User", "1"), json!({"id": "1"}))
613 .await
614 .unwrap();
615 storage
616 .put(StorageKey::new("tenant1", "Group", "1"), json!({"id": "1"}))
617 .await
618 .unwrap();
619 storage
620 .put(StorageKey::new("tenant2", "User", "1"), json!({"id": "1"}))
621 .await
622 .unwrap();
623
624 let mut tenants = storage.list_tenants().await.unwrap();
626 tenants.sort();
627 assert_eq!(tenants, vec!["tenant1", "tenant2"]);
628
629 let mut types1 = storage.list_resource_types("tenant1").await.unwrap();
631 types1.sort();
632 assert_eq!(types1, vec!["Group", "User"]);
633
634 let types2 = storage.list_resource_types("tenant2").await.unwrap();
635 assert_eq!(types2, vec!["User"]);
636
637 let types_none = storage.list_resource_types("nonexistent").await.unwrap();
639 assert!(types_none.is_empty());
640 }
641
642 #[tokio::test]
643 async fn test_extract_attribute_value() {
644 let data = json!({
645 "userName": "john.doe",
646 "emails": [
647 {"value": "john@example.com", "primary": true},
648 {"value": "john.doe@work.com", "primary": false}
649 ],
650 "address": {
651 "street": "123 Main St",
652 "city": "Anytown"
653 }
654 });
655
656 assert_eq!(
658 InMemoryStorage::extract_attribute_value(&data, "userName"),
659 Some("john.doe".to_string())
660 );
661
662 assert_eq!(
664 InMemoryStorage::extract_attribute_value(&data, "address.city"),
665 Some("Anytown".to_string())
666 );
667
668 assert_eq!(
670 InMemoryStorage::extract_attribute_value(&data, "emails.0.value"),
671 Some("john@example.com".to_string())
672 );
673
674 assert_eq!(
676 InMemoryStorage::extract_attribute_value(&data, "emails.0.primary"),
677 Some("true".to_string())
678 );
679
680 assert_eq!(
682 InMemoryStorage::extract_attribute_value(&data, "nonexistent"),
683 None
684 );
685
686 assert_eq!(
688 InMemoryStorage::extract_attribute_value(&data, "emails.99.value"),
689 None
690 );
691 }
692}