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

1""" 

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

3 

4`litellm_accessgrouptable.assigned_key_ids` and `litellm_verificationtoken.access_group_ids` 

5are the two halves of one relationship and BOTH are read: the access group's 

6attached-keys view reads the former, and so does the grant check in 

7`auth_checks.get_authorized_resources_from_key_access_groups`, which authorizes a 

8key only when the group lists the key's token (or the key's team). The access-group 

9endpoints maintain both halves already; this module is what the key write paths call 

10so an edit from that side is mirrored back. 

11 

12Every write is a single guarded statement rather than a read-modify-write. Prisma has no 

13atomic scalar-list removal (see `TeamRepository.remove_member`), and the read-modify-write 

14it otherwise forces is not safe here: a lost update would put an already revoked token back 

15into a group and restore its grants, or drop a grant an admin just made. The guards also 

16make each statement idempotent, so a retry cannot duplicate an entry. Each statement covers 

17every group the request touches at once, so the size of the caller's id list does not turn 

18into a matching number of round trips, and returns the ids it actually moved so only those 

19groups are dropped from cache. 

20 

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

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

23module eagerly from `key_management_endpoints` would put it in `sys.modules` 

24without its router ever being included, which drops its routes from the OpenAPI 

25schema. 

26""" 

27 

28from collections.abc import Sequence 

29from typing import Final, Protocol 

30 

31from pydantic import BaseModel 

32 

33from litellm.proxy._types import ( 

34 LiteLLM_VerificationToken, 

35 RegenerateKeyRequest, 

36 UpdateKeyRequest, 

37) 

38from litellm.proxy.auth.auth_checks import ( 

39 _delete_cache_access_object, # pyright: ignore[reportPrivateUsage] # the access-group endpoints reach for this same cache primitive 

40) 

41from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper 

42from litellm.repositories.table_repositories import AccessGroupRepository 

43 

44 

45class _MovedGroupRow(BaseModel): 

46 access_group_id: str 

47 

48 

49class _RawExecutor(Protocol): 

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

51 

52 

53_ATTACH_KEY_SQL: Final = ( 

54 'UPDATE "LiteLLM_AccessGroupTable" ' 

55 'SET "assigned_key_ids" = array_append("assigned_key_ids", $1) ' 

56 'WHERE "access_group_id" = ANY($2::text[]) AND NOT ($1 = ANY("assigned_key_ids")) ' 

57 'RETURNING "access_group_id"' 

58) 

59 

60_DETACH_KEY_SQL: Final = ( 

61 'UPDATE "LiteLLM_AccessGroupTable" ' 

62 'SET "assigned_key_ids" = array_remove("assigned_key_ids", $1) ' 

63 'WHERE "access_group_id" = ANY($2::text[]) AND $1 = ANY("assigned_key_ids") ' 

64 'RETURNING "access_group_id"' 

65) 

66 

67_REPOINT_KEY_SQL: Final = ( 

68 'UPDATE "LiteLLM_AccessGroupTable" ' 

69 'SET "assigned_key_ids" = array_append(array_remove(array_remove("assigned_key_ids", $1), $2), $2) ' 

70 'WHERE $1 = ANY("assigned_key_ids") ' 

71 'RETURNING "access_group_id"' 

72) 

73 

74 

75def _raw_executor(prisma_client: object) -> _RawExecutor: 

76 """Narrow the untyped Prisma client down to the raw-query call this module makes, pinned to the writer.""" 

77 db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client 

78 return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin 

79 

80 

81async def _invalidate_access_group_cache(access_group_id: str) -> None: 

82 """ 

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

84 

85 Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server 

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

87 """ 

88 from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache 

89 

90 await _delete_cache_access_object( 

91 access_group_id=access_group_id, 

92 user_api_key_cache=user_api_key_cache, 

93 proxy_logging_obj=proxy_logging_obj, 

94 ) 

95 

96 

97async def _invalidate_moved_groups(moved_rows: Sequence[object]) -> None: 

98 for row in moved_rows: 98 ↛ 99line 98 didn't jump to line 99 because the loop on line 98 never started

99 await _invalidate_access_group_cache(_MovedGroupRow.model_validate(row).access_group_id) 

100 

101 

102async def _write_membership(prisma_client: object, sql: str, access_group_ids: frozenset[str], key_token: str) -> None: 

103 """Run one guarded membership statement for every listed group, dropping the cache of those it moved.""" 

104 if not access_group_ids: 

105 return 

106 await _invalidate_moved_groups( 

107 await _raw_executor(prisma_client).query_raw(sql, key_token, sorted(access_group_ids)) 

108 ) 

109 

110 

111async def sync_key_access_group_membership( 

112 prisma_client: object, 

113 key_token: str, 

114 previous_access_group_ids: Sequence[str] | None, 

115 updated_access_group_ids: Sequence[str] | None, 

116) -> None: 

117 """Mirror a key-side change to `access_group_ids` onto each access group's `assigned_key_ids`.""" 

118 previous: Final = frozenset(previous_access_group_ids or ()) 

119 updated: Final = frozenset(updated_access_group_ids or ()) 

120 

121 await _write_membership(prisma_client, _ATTACH_KEY_SQL, updated - previous, key_token) 

122 await _write_membership(prisma_client, _DETACH_KEY_SQL, previous - updated, key_token) 

123 

124 

125async def sync_key_update_access_group_membership( 

126 prisma_client: object, 

127 key_token: str, 

128 data: UpdateKeyRequest | RegenerateKeyRequest, 

129 existing_key_row: LiteLLM_VerificationToken, 

130) -> None: 

131 """ 

132 Mirror a key UPDATE onto the group side, honouring `exclude_unset` semantics. 

133 

134 The key row is written from `model_dump(exclude_unset=True)`, so a request that never 

135 mentions `access_group_ids` leaves the key's own list alone and must leave the group's 

136 copy alone too. Reading the attribute instead of `model_fields_set` would see None on 

137 every unrelated edit and withdraw the token from every group it belongs to. 

138 """ 

139 if "access_group_ids" not in data.model_fields_set: 

140 return 

141 await sync_key_access_group_membership( 

142 prisma_client=prisma_client, 

143 key_token=key_token, 

144 previous_access_group_ids=existing_key_row.access_group_ids, 

145 updated_access_group_ids=data.access_group_ids, 

146 ) 

147 

148 

149async def sync_key_regeneration_access_group_membership( 

150 prisma_client: object, 

151 previous_key_token: str, 

152 new_key_token: str, 

153 data: RegenerateKeyRequest | None, 

154 existing_key_row: LiteLLM_VerificationToken, 

155) -> None: 

156 """ 

157 Re-point every group's copy from the old token to the regenerated one. 

158 

159 Regeneration replaces the token, which is the identity `assigned_key_ids` stores, so 

160 leaving the old hash behind both points the group at a row that no longer exists and 

161 denies the regenerated key the group's grants. The swap is driven by the groups that 

162 hold the old token when the statement runs, not by the key row read earlier, so a group 

163 edited in between is neither resurrected nor skipped. Removing the new token before 

164 appending it keeps a re-run from duplicating it. 

165 """ 

166 await _invalidate_moved_groups( 

167 await _raw_executor(prisma_client).query_raw(_REPOINT_KEY_SQL, previous_key_token, new_key_token) 

168 ) 

169 if data is not None: 

170 await sync_key_update_access_group_membership( 

171 prisma_client=prisma_client, 

172 key_token=new_key_token, 

173 data=data, 

174 existing_key_row=existing_key_row, 

175 )