Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_helpers/access_group_team_sync.py: 83%

1""" 

2Reverse sync for the team side of the team <-> access group relationship. 

3 

4`litellm_accessgrouptable.assigned_team_ids` and `litellm_teamtable.access_group_ids` 

5are two copies of the same relationship, and both are read: the access group's 

6attached-teams view reads the former, and so does the key-side grant check in 

7`auth_checks.get_authorized_resources_from_key_access_groups`. The access-group 

8endpoints maintain both copies already; this module is what the team write paths 

9call so an edit from that side is mirrored back. 

10 

11It deliberately lives outside `access_group_endpoints`, which is a lazily 

12registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that 

13module eagerly from `team_endpoints` would put it in `sys.modules` without its 

14router ever being included, which drops its routes from the OpenAPI schema. 

15""" 

16 

17import asyncio 

18from collections.abc import Mapping, Sequence 

19from typing import Final, Protocol 

20 

21from pydantic import BaseModel, TypeAdapter 

22 

23from litellm.proxy.auth.auth_checks import _delete_cache_access_object 

24 

25# hashtext collisions only cost two unrelated teams a little serialization, and the 

26# lock is never taken by the access-group endpoints as a SELECT ... FOR UPDATE row lock, 

27# so it cannot join their access-group-then-team lock order to form a cycle. team_endpoints 

28# reuses this exact statement to serialize /team/member_add and /team/delete against each 

29# other and against this mirror, rather than defining a second, divergent lock on the same key. 

30TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" 

31 

32_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1' 

33 

34# The groups the team is on either side of the reconcile, so the cache step is driven by 

35# desired state rather than by which rows this attempt happened to change. A retry after a 

36# failed invalidation finds the same set even though its statements are already no-ops. 

37_AFFECTED_SQL: Final = """ 

38SELECT access_group_id FROM "LiteLLM_AccessGroupTable" 

39WHERE access_group_id = ANY($2::TEXT[]) 

40 OR $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])) 

41""" 

42 

43_ATTACH_SQL: Final = """ 

44UPDATE "LiteLLM_AccessGroupTable" 

45SET assigned_team_ids = array_append(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]), $1) 

46WHERE access_group_id = ANY($2::TEXT[]) 

47 AND NOT ($1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))) 

48RETURNING access_group_id 

49""" 

50 

51_DETACH_SQL: Final = """ 

52UPDATE "LiteLLM_AccessGroupTable" 

53SET assigned_team_ids = array_remove(assigned_team_ids, $1) 

54WHERE $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])) 

55 AND NOT (access_group_id = ANY($2::TEXT[])) 

56RETURNING access_group_id 

57""" 

58 

59 

60class _AffectedGroup(BaseModel): 

61 access_group_id: str 

62 

63 

64class _TeamGroups(BaseModel): 

65 access_group_ids: tuple[str, ...] | None = None 

66 

67 

68_AffectedGroups: Final = TypeAdapter(tuple[_AffectedGroup, ...]) 

69_TeamRows: Final = TypeAdapter(tuple[_TeamGroups, ...]) 

70 

71 

72class AccessGroupSyncTx(Protocol): 

73 async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... 73 ↛ exitline 73 didn't return from function 'query_raw' because

74 

75 

76class _Transaction(Protocol): 

77 async def __aenter__(self) -> AccessGroupSyncTx: ... 77 ↛ exitline 77 didn't return from function '__aenter__' because

78 

79 async def __aexit__(self, *exc_info: object) -> None: ... 79 ↛ exitline 79 didn't return from function '__aexit__' because

80 

81 

82class _PrismaDb(Protocol): 

83 def tx(self) -> _Transaction: ... 83 ↛ exitline 83 didn't return from function 'tx' because

84 

85 

86class _PrismaClient(Protocol): 

87 @property 

88 def db(self) -> _PrismaDb: ... 88 ↛ exitline 88 didn't return from function 'db' because

89 

90 

91async def invalidate_access_group_cache(access_group_id: str) -> None: 

92 """ 

93 Drop an access group entry from both the in-memory and Redis caches. 

94 

95 Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server 

96 to avoid circular imports, following the same pattern as key_management_endpoints. 

97 """ 

98 from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache 

99 

100 await _delete_cache_access_object( 

101 access_group_id=access_group_id, 

102 user_api_key_cache=user_api_key_cache, 

103 proxy_logging_obj=proxy_logging_obj, 

104 ) 

105 

106 

107async def invalidate_access_group_caches(access_group_ids: Sequence[str]) -> None: 

108 """ 

109 Drop every given access group from the caches, then raise if any drop failed. 

110 

111 Every entry is attempted even when one raises, so a single unreachable cache cannot 

112 leave the rest of the reconciled groups serving a grant the admin revoked. 

113 """ 

114 outcomes: Final = await asyncio.gather( 

115 *(invalidate_access_group_cache(access_group_id) for access_group_id in access_group_ids), 

116 return_exceptions=True, 

117 ) 

118 for outcome in outcomes: 118 ↛ 119line 118 didn't jump to line 119 because the loop on line 118 never started

119 if isinstance(outcome, BaseException): 

120 raise outcome 

121 

122 

123async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id: str) -> tuple[str, ...]: 

124 """ 

125 Reconcile every access group's `assigned_team_ids` against the team's own 

126 `access_group_ids`, and return the groups whose cache the caller has to drop once the 

127 transaction commits. 

128 

129 Call this inside the transaction that writes the team row, or after that row is 

130 written or deleted: a team with no row reconciles to an empty set, which detaches it 

131 from every group. 

132 

133 The team row is read here rather than passed in, under an advisory lock held for the 

134 rest of the transaction. That is what makes concurrent writes to the same team 

135 converge, since each mirror reconciles against the row as the transaction sees it 

136 instead of against the snapshot its own caller happened to see. It also means a retry 

137 heals a sync that failed partway, where a before/after delta would compute nothing. 

138 

139 Both mirror statements are set-based and mutate the array inside the statement, so a 

140 concurrent write for a different team cannot be lost the way a read-modify-write of 

141 the whole array can, and the pair commits together or not at all. 

142 """ 

143 await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) 

144 team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id)) 

145 desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else () 

146 affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired)) 

147 await tx.query_raw(_ATTACH_SQL, team_id, desired) 

148 await tx.query_raw(_DETACH_SQL, team_id, desired) 

149 return tuple(group.access_group_id for group in affected) 

150 

151 

152async def sync_team_access_group_membership(prisma_client: _PrismaClient, team_id: str) -> None: 

153 """Reconcile the mirror for an already committed team write, in its own transaction.""" 

154 async with prisma_client.db.tx() as tx: 

155 affected: Final = await reconcile_team_access_group_membership(tx, team_id) 

156 

157 await invalidate_access_group_caches(affected)