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