litellm_rs/core/teams/
repository.rs

1//! Team repository trait and implementations
2//!
3//! This module provides storage abstractions for team management.
4
5use crate::core::models::team::{Team, TeamMember, TeamRole, TeamStatus};
6use crate::utils::error::gateway_error::{GatewayError, Result};
7use async_trait::async_trait;
8use std::collections::HashMap;
9use std::sync::RwLock;
10use uuid::Uuid;
11
12/// Team repository trait for abstracting storage operations
13#[async_trait]
14pub trait TeamRepository: Send + Sync {
15    /// Create a new team
16    async fn create(&self, team: Team) -> Result<Team>;
17
18    /// Get a team by ID
19    async fn get(&self, id: Uuid) -> Result<Option<Team>>;
20
21    /// Get a team by name
22    async fn get_by_name(&self, name: &str) -> Result<Option<Team>>;
23
24    /// Update an existing team
25    async fn update(&self, team: Team) -> Result<Team>;
26
27    /// Delete a team by ID
28    async fn delete(&self, id: Uuid) -> Result<()>;
29
30    /// List all teams with pagination
31    async fn list(&self, offset: u32, limit: u32) -> Result<(Vec<Team>, u64)>;
32
33    /// Count total teams
34    async fn count(&self) -> Result<u64>;
35
36    /// Add a member to a team
37    async fn add_member(&self, member: TeamMember) -> Result<TeamMember>;
38
39    /// Get a team member
40    async fn get_member(&self, team_id: Uuid, user_id: Uuid) -> Result<Option<TeamMember>>;
41
42    /// Update a team member's role
43    async fn update_member_role(
44        &self,
45        team_id: Uuid,
46        user_id: Uuid,
47        role: TeamRole,
48    ) -> Result<TeamMember>;
49
50    /// Remove a member from a team
51    async fn remove_member(&self, team_id: Uuid, user_id: Uuid) -> Result<()>;
52
53    /// List members of a team
54    async fn list_members(&self, team_id: Uuid) -> Result<Vec<TeamMember>>;
55
56    /// Get teams for a user
57    async fn get_user_teams(&self, user_id: Uuid) -> Result<Vec<Team>>;
58}
59
60/// In-memory team repository for testing and development
61pub struct InMemoryTeamRepository {
62    teams: RwLock<HashMap<Uuid, Team>>,
63    members: RwLock<HashMap<(Uuid, Uuid), TeamMember>>,
64}
65
66impl InMemoryTeamRepository {
67    /// Create a new in-memory repository
68    pub fn new() -> Self {
69        Self {
70            teams: RwLock::new(HashMap::new()),
71            members: RwLock::new(HashMap::new()),
72        }
73    }
74}
75
76impl Default for InMemoryTeamRepository {
77    fn default() -> Self {
78        Self::new()
79    }
80}
81
82#[async_trait]
83impl TeamRepository for InMemoryTeamRepository {
84    async fn create(&self, team: Team) -> Result<Team> {
85        let mut teams = self
86            .teams
87            .write()
88            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
89
90        // Check for name conflict
91        if teams.values().any(|t| t.name == team.name) {
92            return Err(GatewayError::Conflict(format!(
93                "Team with name '{}' already exists",
94                team.name
95            )));
96        }
97
98        let id = team.id();
99        teams.insert(id, team.clone());
100        Ok(team)
101    }
102
103    async fn get(&self, id: Uuid) -> Result<Option<Team>> {
104        let teams = self
105            .teams
106            .read()
107            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
108        Ok(teams.get(&id).cloned())
109    }
110
111    async fn get_by_name(&self, name: &str) -> Result<Option<Team>> {
112        let teams = self
113            .teams
114            .read()
115            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
116        Ok(teams.values().find(|t| t.name == name).cloned())
117    }
118
119    async fn update(&self, mut team: Team) -> Result<Team> {
120        let mut teams = self
121            .teams
122            .write()
123            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
124
125        let id = team.id();
126        if !teams.contains_key(&id) {
127            return Err(GatewayError::NotFound(format!("Team {} not found", id)));
128        }
129
130        // Check for name conflict with other teams
131        if teams.values().any(|t| t.name == team.name && t.id() != id) {
132            return Err(GatewayError::Conflict(format!(
133                "Team with name '{}' already exists",
134                team.name
135            )));
136        }
137
138        team.metadata.touch();
139        teams.insert(id, team.clone());
140        Ok(team)
141    }
142
143    async fn delete(&self, id: Uuid) -> Result<()> {
144        let mut teams = self
145            .teams
146            .write()
147            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
148
149        if teams.remove(&id).is_none() {
150            return Err(GatewayError::NotFound(format!("Team {} not found", id)));
151        }
152
153        // Remove all members of the team
154        let mut members = self
155            .members
156            .write()
157            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
158        members.retain(|(team_id, _), _| *team_id != id);
159
160        Ok(())
161    }
162
163    async fn list(&self, offset: u32, limit: u32) -> Result<(Vec<Team>, u64)> {
164        let teams = self
165            .teams
166            .read()
167            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
168
169        let total = teams.len() as u64;
170        let mut team_list: Vec<Team> = teams
171            .values()
172            .filter(|t| !matches!(t.status, TeamStatus::Deleted))
173            .cloned()
174            .collect();
175
176        // Sort by created_at descending
177        team_list.sort_by_key(|b| std::cmp::Reverse(b.metadata.created_at));
178
179        let paginated: Vec<Team> = team_list
180            .into_iter()
181            .skip(offset as usize)
182            .take(limit as usize)
183            .collect();
184
185        Ok((paginated, total))
186    }
187
188    async fn count(&self) -> Result<u64> {
189        let teams = self
190            .teams
191            .read()
192            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
193        Ok(teams
194            .values()
195            .filter(|t| !matches!(t.status, TeamStatus::Deleted))
196            .count() as u64)
197    }
198
199    async fn add_member(&self, member: TeamMember) -> Result<TeamMember> {
200        // Verify team exists
201        {
202            let teams = self
203                .teams
204                .read()
205                .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
206            if !teams.contains_key(&member.team_id) {
207                return Err(GatewayError::NotFound(format!(
208                    "Team {} not found",
209                    member.team_id
210                )));
211            }
212        }
213
214        let mut members = self
215            .members
216            .write()
217            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
218
219        let key = (member.team_id, member.user_id);
220        if members.contains_key(&key) {
221            return Err(GatewayError::Conflict(format!(
222                "User {} is already a member of team {}",
223                member.user_id, member.team_id
224            )));
225        }
226
227        members.insert(key, member.clone());
228        Ok(member)
229    }
230
231    async fn get_member(&self, team_id: Uuid, user_id: Uuid) -> Result<Option<TeamMember>> {
232        let members = self
233            .members
234            .read()
235            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
236        Ok(members.get(&(team_id, user_id)).cloned())
237    }
238
239    async fn update_member_role(
240        &self,
241        team_id: Uuid,
242        user_id: Uuid,
243        role: TeamRole,
244    ) -> Result<TeamMember> {
245        let mut members = self
246            .members
247            .write()
248            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
249
250        let key = (team_id, user_id);
251        let member = members.get_mut(&key).ok_or_else(|| {
252            GatewayError::NotFound(format!("Member {} not found in team {}", user_id, team_id))
253        })?;
254
255        member.role = role;
256        member.metadata.touch();
257        Ok(member.clone())
258    }
259
260    async fn remove_member(&self, team_id: Uuid, user_id: Uuid) -> Result<()> {
261        let mut members = self
262            .members
263            .write()
264            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
265
266        let key = (team_id, user_id);
267        if members.remove(&key).is_none() {
268            return Err(GatewayError::NotFound(format!(
269                "Member {} not found in team {}",
270                user_id, team_id
271            )));
272        }
273
274        Ok(())
275    }
276
277    async fn list_members(&self, team_id: Uuid) -> Result<Vec<TeamMember>> {
278        let members = self
279            .members
280            .read()
281            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
282
283        let team_members: Vec<TeamMember> = members
284            .iter()
285            .filter(|((tid, _), _)| *tid == team_id)
286            .map(|(_, m)| m.clone())
287            .collect();
288
289        Ok(team_members)
290    }
291
292    async fn get_user_teams(&self, user_id: Uuid) -> Result<Vec<Team>> {
293        let members = self
294            .members
295            .read()
296            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
297
298        let team_ids: Vec<Uuid> = members
299            .iter()
300            .filter(|((_, uid), _)| *uid == user_id)
301            .map(|((tid, _), _)| *tid)
302            .collect();
303
304        let teams = self
305            .teams
306            .read()
307            .map_err(|_| GatewayError::Internal("Lock poisoned".to_string()))?;
308
309        let user_teams: Vec<Team> = team_ids
310            .into_iter()
311            .filter_map(|tid| teams.get(&tid).cloned())
312            .collect();
313
314        Ok(user_teams)
315    }
316}
317
318#[cfg(feature = "postgres")]
319pub mod postgres {
320    //! PostgreSQL implementation of TeamRepository
321    //!
322    //! This module provides a PostgreSQL-backed implementation.
323
324    use super::*;
325    use crate::storage::database::Database;
326    use std::sync::Arc;
327
328    /// PostgreSQL team repository
329    pub struct PostgresTeamRepository {
330        _db: Arc<Database>,
331    }
332
333    impl PostgresTeamRepository {
334        /// Create a new PostgreSQL repository
335        pub fn new(db: Arc<Database>) -> Self {
336            Self { _db: db }
337        }
338    }
339
340    #[async_trait]
341    impl TeamRepository for PostgresTeamRepository {
342        async fn create(&self, _team: Team) -> Result<Team> {
343            // NOTE: PostgreSQL storage for teams not yet implemented.
344            // This is a placeholder that would use SeaORM entities
345            Err(GatewayError::NotImplemented(
346                "PostgreSQL team storage not yet implemented".to_string(),
347            ))
348        }
349
350        async fn get(&self, _id: Uuid) -> Result<Option<Team>> {
351            Err(GatewayError::NotImplemented(
352                "PostgreSQL team storage not yet implemented".to_string(),
353            ))
354        }
355
356        async fn get_by_name(&self, _name: &str) -> Result<Option<Team>> {
357            Err(GatewayError::NotImplemented(
358                "PostgreSQL team storage not yet implemented".to_string(),
359            ))
360        }
361
362        async fn update(&self, _team: Team) -> Result<Team> {
363            Err(GatewayError::NotImplemented(
364                "PostgreSQL team storage not yet implemented".to_string(),
365            ))
366        }
367
368        async fn delete(&self, _id: Uuid) -> Result<()> {
369            Err(GatewayError::NotImplemented(
370                "PostgreSQL team storage not yet implemented".to_string(),
371            ))
372        }
373
374        async fn list(&self, _offset: u32, _limit: u32) -> Result<(Vec<Team>, u64)> {
375            Err(GatewayError::NotImplemented(
376                "PostgreSQL team storage not yet implemented".to_string(),
377            ))
378        }
379
380        async fn count(&self) -> Result<u64> {
381            Err(GatewayError::NotImplemented(
382                "PostgreSQL team storage not yet implemented".to_string(),
383            ))
384        }
385
386        async fn add_member(&self, _member: TeamMember) -> Result<TeamMember> {
387            Err(GatewayError::NotImplemented(
388                "PostgreSQL team storage not yet implemented".to_string(),
389            ))
390        }
391
392        async fn get_member(&self, _team_id: Uuid, _user_id: Uuid) -> Result<Option<TeamMember>> {
393            Err(GatewayError::NotImplemented(
394                "PostgreSQL team storage not yet implemented".to_string(),
395            ))
396        }
397
398        async fn update_member_role(
399            &self,
400            _team_id: Uuid,
401            _user_id: Uuid,
402            _role: TeamRole,
403        ) -> Result<TeamMember> {
404            Err(GatewayError::NotImplemented(
405                "PostgreSQL team storage not yet implemented".to_string(),
406            ))
407        }
408
409        async fn remove_member(&self, _team_id: Uuid, _user_id: Uuid) -> Result<()> {
410            Err(GatewayError::NotImplemented(
411                "PostgreSQL team storage not yet implemented".to_string(),
412            ))
413        }
414
415        async fn list_members(&self, _team_id: Uuid) -> Result<Vec<TeamMember>> {
416            Err(GatewayError::NotImplemented(
417                "PostgreSQL team storage not yet implemented".to_string(),
418            ))
419        }
420
421        async fn get_user_teams(&self, _user_id: Uuid) -> Result<Vec<Team>> {
422            Err(GatewayError::NotImplemented(
423                "PostgreSQL team storage not yet implemented".to_string(),
424            ))
425        }
426    }
427}
428
429#[cfg(test)]
430mod tests {
431    use super::*;
432
433    #[tokio::test]
434    async fn test_create_team() {
435        let repo = InMemoryTeamRepository::new();
436        let team = Team::new("test-team".to_string(), Some("Test Team".to_string()));
437
438        let created = repo.create(team.clone()).await.unwrap();
439        assert_eq!(created.name, "test-team");
440    }
441
442    #[tokio::test]
443    async fn test_get_team() {
444        let repo = InMemoryTeamRepository::new();
445        let team = Team::new("test-team".to_string(), None);
446        let id = team.id();
447
448        repo.create(team).await.unwrap();
449        let fetched = repo.get(id).await.unwrap();
450
451        assert!(fetched.is_some());
452        assert_eq!(fetched.unwrap().name, "test-team");
453    }
454
455    #[tokio::test]
456    async fn test_get_team_not_found() {
457        let repo = InMemoryTeamRepository::new();
458        let result = repo.get(Uuid::new_v4()).await.unwrap();
459        assert!(result.is_none());
460    }
461
462    #[tokio::test]
463    async fn test_get_by_name() {
464        let repo = InMemoryTeamRepository::new();
465        let team = Team::new("unique-name".to_string(), None);
466        repo.create(team).await.unwrap();
467
468        let fetched = repo.get_by_name("unique-name").await.unwrap();
469        assert!(fetched.is_some());
470        assert_eq!(fetched.unwrap().name, "unique-name");
471    }
472
473    #[tokio::test]
474    async fn test_update_team() {
475        let repo = InMemoryTeamRepository::new();
476        let team = Team::new("original".to_string(), None);
477        let id = team.id();
478
479        repo.create(team).await.unwrap();
480
481        let mut updated_team = repo.get(id).await.unwrap().unwrap();
482        updated_team.description = Some("Updated description".to_string());
483
484        let result = repo.update(updated_team).await.unwrap();
485        assert_eq!(result.description, Some("Updated description".to_string()));
486    }
487
488    #[tokio::test]
489    async fn test_delete_team() {
490        let repo = InMemoryTeamRepository::new();
491        let team = Team::new("to-delete".to_string(), None);
492        let id = team.id();
493
494        repo.create(team).await.unwrap();
495        repo.delete(id).await.unwrap();
496
497        let fetched = repo.get(id).await.unwrap();
498        assert!(fetched.is_none());
499    }
500
501    #[tokio::test]
502    async fn test_list_teams() {
503        let repo = InMemoryTeamRepository::new();
504
505        for i in 0..5 {
506            let team = Team::new(format!("team-{}", i), None);
507            repo.create(team).await.unwrap();
508        }
509
510        let (teams, total) = repo.list(0, 10).await.unwrap();
511        assert_eq!(teams.len(), 5);
512        assert_eq!(total, 5);
513    }
514
515    #[tokio::test]
516    async fn test_list_teams_pagination() {
517        let repo = InMemoryTeamRepository::new();
518
519        for i in 0..10 {
520            let team = Team::new(format!("team-{}", i), None);
521            repo.create(team).await.unwrap();
522        }
523
524        let (teams, total) = repo.list(0, 5).await.unwrap();
525        assert_eq!(teams.len(), 5);
526        assert_eq!(total, 10);
527
528        let (teams2, _) = repo.list(5, 5).await.unwrap();
529        assert_eq!(teams2.len(), 5);
530    }
531
532    #[tokio::test]
533    async fn test_duplicate_team_name() {
534        let repo = InMemoryTeamRepository::new();
535        let team1 = Team::new("duplicate".to_string(), None);
536        let team2 = Team::new("duplicate".to_string(), None);
537
538        repo.create(team1).await.unwrap();
539        let result = repo.create(team2).await;
540
541        assert!(result.is_err());
542    }
543
544    #[tokio::test]
545    async fn test_add_member() {
546        let repo = InMemoryTeamRepository::new();
547        let team = Team::new("team".to_string(), None);
548        let team_id = team.id();
549        repo.create(team).await.unwrap();
550
551        let user_id = Uuid::new_v4();
552        let member = TeamMember::new(team_id, user_id, TeamRole::Member, None);
553
554        let created = repo.add_member(member).await.unwrap();
555        assert_eq!(created.user_id, user_id);
556        assert!(matches!(created.role, TeamRole::Member));
557    }
558
559    #[tokio::test]
560    async fn test_get_member() {
561        let repo = InMemoryTeamRepository::new();
562        let team = Team::new("team".to_string(), None);
563        let team_id = team.id();
564        repo.create(team).await.unwrap();
565
566        let user_id = Uuid::new_v4();
567        let member = TeamMember::new(team_id, user_id, TeamRole::Admin, None);
568        repo.add_member(member).await.unwrap();
569
570        let fetched = repo.get_member(team_id, user_id).await.unwrap();
571        assert!(fetched.is_some());
572        assert!(matches!(fetched.unwrap().role, TeamRole::Admin));
573    }
574
575    #[tokio::test]
576    async fn test_update_member_role() {
577        let repo = InMemoryTeamRepository::new();
578        let team = Team::new("team".to_string(), None);
579        let team_id = team.id();
580        repo.create(team).await.unwrap();
581
582        let user_id = Uuid::new_v4();
583        let member = TeamMember::new(team_id, user_id, TeamRole::Member, None);
584        repo.add_member(member).await.unwrap();
585
586        let updated = repo
587            .update_member_role(team_id, user_id, TeamRole::Admin)
588            .await
589            .unwrap();
590        assert!(matches!(updated.role, TeamRole::Admin));
591    }
592
593    #[tokio::test]
594    async fn test_remove_member() {
595        let repo = InMemoryTeamRepository::new();
596        let team = Team::new("team".to_string(), None);
597        let team_id = team.id();
598        repo.create(team).await.unwrap();
599
600        let user_id = Uuid::new_v4();
601        let member = TeamMember::new(team_id, user_id, TeamRole::Member, None);
602        repo.add_member(member).await.unwrap();
603
604        repo.remove_member(team_id, user_id).await.unwrap();
605
606        let fetched = repo.get_member(team_id, user_id).await.unwrap();
607        assert!(fetched.is_none());
608    }
609
610    #[tokio::test]
611    async fn test_list_members() {
612        let repo = InMemoryTeamRepository::new();
613        let team = Team::new("team".to_string(), None);
614        let team_id = team.id();
615        repo.create(team).await.unwrap();
616
617        for _ in 0..3 {
618            let member = TeamMember::new(team_id, Uuid::new_v4(), TeamRole::Member, None);
619            repo.add_member(member).await.unwrap();
620        }
621
622        let members = repo.list_members(team_id).await.unwrap();
623        assert_eq!(members.len(), 3);
624    }
625
626    #[tokio::test]
627    async fn test_get_user_teams() {
628        let repo = InMemoryTeamRepository::new();
629        let user_id = Uuid::new_v4();
630
631        for i in 0..3 {
632            let team = Team::new(format!("team-{}", i), None);
633            let team_id = team.id();
634            repo.create(team).await.unwrap();
635
636            let member = TeamMember::new(team_id, user_id, TeamRole::Member, None);
637            repo.add_member(member).await.unwrap();
638        }
639
640        let teams = repo.get_user_teams(user_id).await.unwrap();
641        assert_eq!(teams.len(), 3);
642    }
643
644    #[tokio::test]
645    async fn test_add_member_to_nonexistent_team() {
646        let repo = InMemoryTeamRepository::new();
647        let member = TeamMember::new(Uuid::new_v4(), Uuid::new_v4(), TeamRole::Member, None);
648
649        let result = repo.add_member(member).await;
650        assert!(result.is_err());
651    }
652
653    #[tokio::test]
654    async fn test_duplicate_member() {
655        let repo = InMemoryTeamRepository::new();
656        let team = Team::new("team".to_string(), None);
657        let team_id = team.id();
658        repo.create(team).await.unwrap();
659
660        let user_id = Uuid::new_v4();
661        let member1 = TeamMember::new(team_id, user_id, TeamRole::Member, None);
662        let member2 = TeamMember::new(team_id, user_id, TeamRole::Admin, None);
663
664        repo.add_member(member1).await.unwrap();
665        let result = repo.add_member(member2).await;
666
667        assert!(result.is_err());
668    }
669}
litellm_rs/core/teams/repository.rs

litellm_rs/core/teams/
repository.rs
