Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/master_key_migration.py: 32%

1import json 

2from collections.abc import Awaitable, Callable, Mapping 

3from dataclasses import dataclass 

4from enum import Enum 

5from typing import Final 

6 

7from pydantic import JsonValue, TypeAdapter 

8from typing_extensions import assert_never 

9 

10from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH 

11from litellm.proxy.auth.master_key_boot_check import MIGRATE_FROM_MASTER_KEY_ENV_VAR, SALT_KEY_ENV_VAR 

12from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX 

13from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_value_helper 

14from litellm.proxy.db.create_views import SupportsRawQueries 

15 

16 

17@dataclass(frozen=True, slots=True) 

18class _SecretColumn: 

19 table: str 

20 primary_key: str 

21 column: str 

22 is_json: bool = True 

23 only_rows_with_marked_ciphertexts: bool = False 

24 

25 

26_SECRET_COLUMNS: Final = ( 

27 _SecretColumn("LiteLLM_ProxyModelTable", "model_id", "litellm_params"), 

28 _SecretColumn("LiteLLM_CredentialsTable", "credential_id", "credential_values"), 

29 _SecretColumn("LiteLLM_Config", "param_name", "param_value"), 

30 _SecretColumn("LiteLLM_SSOConfig", "id", "sso_settings"), 

31 _SecretColumn("LiteLLM_CacheConfig", "id", "cache_settings"), 

32 _SecretColumn("LiteLLM_ConfigOverrides", "config_type", "config_value"), 

33 _SecretColumn("LiteLLM_MCPServerTable", "server_id", "credentials"), 

34 _SecretColumn("LiteLLM_MCPServerTable", "server_id", "static_headers"), 

35 _SecretColumn("LiteLLM_MCPServerTable", "server_id", "env_vars"), 

36 _SecretColumn("LiteLLM_MCPServerTable", "server_id", "env"), 

37 _SecretColumn("LiteLLM_MCPServerOAuthClient", "server_id", "credentials"), 

38 _SecretColumn("LiteLLM_MCPUserCredentials", "id", "credential_b64", is_json=False), 

39 _SecretColumn("LiteLLM_MCPUserEnvVars", "id", "values_b64", is_json=False), 

40 _SecretColumn("LiteLLM_SSOIdentityAssertion", "user_id", "assertion_b64", is_json=False), 

41 _SecretColumn("LiteLLM_TeamTable", "team_id", "metadata", only_rows_with_marked_ciphertexts=True), 

42 _SecretColumn("LiteLLM_VerificationToken", "token", "metadata", only_rows_with_marked_ciphertexts=True), 

43 _SecretColumn("LiteLLM_UserTable", "user_id", "metadata", only_rows_with_marked_ciphertexts=True), 

44 _SecretColumn("LiteLLM_DeletedTeamTable", "id", "metadata", only_rows_with_marked_ciphertexts=True), 

45 _SecretColumn("LiteLLM_DeletedVerificationToken", "id", "metadata", only_rows_with_marked_ciphertexts=True), 

46) 

47 

48_STORED_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) 

49_PRIMARY_KEY: Final = TypeAdapter(str) 

50 

51ReplaceCiphertext = Callable[[str], str | None] 

52 

53 

54def replace_ciphertexts(value: JsonValue, replacement_for: ReplaceCiphertext, depth: int = 0) -> tuple[JsonValue, int]: 

55 if depth > DEFAULT_MAX_RECURSE_DEPTH: 

56 return value, 0 

57 match value: 

58 case str(): 

59 marker: Final = CALLBACK_VAR_ENCRYPTED_PREFIX if value.startswith(CALLBACK_VAR_ENCRYPTED_PREFIX) else "" 

60 replacement: Final = replacement_for(value.removeprefix(marker)) 

61 return (value, 0) if replacement is None else (marker + replacement, 1) 

62 case list(): 

63 items: Final = tuple(replace_ciphertexts(item, replacement_for, depth + 1) for item in value) 

64 return [item for item, _ in items], sum(count for _, count in items) 

65 case dict(): 

66 fields: Final = {key: replace_ciphertexts(item, replacement_for, depth + 1) for key, item in value.items()} 

67 return {key: item for key, (item, _) in fields.items()}, sum(count for _, count in fields.values()) 

68 case _: 

69 return value, 0 

70 

71 

72async def count_values_encrypted_with(database: SupportsRawQueries, signing_key: str) -> int: 

73 def keep(value: str) -> str | None: 

74 return None if decrypt_if_encrypted_with(value, signing_key) is None else value 

75 

76 return sum( 

77 [ 

78 replace_ciphertexts(_STORED_VALUE.validate_python(row[secret_column.column]), keep)[1] 

79 for secret_column in await _secret_columns_in(database) 

80 for row in await _rows_of(database, secret_column) 

81 ] 

82 ) 

83 

84 

85async def count_values_encrypted_with_or_none( 

86 connect: Callable[[], Awaitable[SupportsRawQueries]], signing_key: str 

87) -> int | None: 

88 try: 

89 return await count_values_encrypted_with(await connect(), signing_key) 

90 except Exception: # noqa: BLE001 # an unreadable database must not replace the boot refusal with a traceback 

91 return None 

92 

93 

94async def reencrypt_stored_values(database: SupportsRawQueries, *, from_key: str, to_key: str) -> int: 

95 def reencrypted(value: str) -> str | None: 

96 plaintext: Final = decrypt_if_encrypted_with(value, from_key) 

97 return None if plaintext is None else _CIPHERTEXT.validate_python(encrypt_value_helper(plaintext, to_key)) 

98 

99 return sum( 

100 [ 

101 await _reencrypt_row(database, secret_column, row, reencrypted) 

102 for secret_column in await _secret_columns_in(database) 

103 for row in await _rows_of(database, secret_column) 

104 ] 

105 ) 

106 

107 

108_CIPHERTEXT: Final = TypeAdapter(str) 

109 

110 

111async def _secret_columns_in(database: SupportsRawQueries) -> tuple[_SecretColumn, ...]: 

112 existing: Final = frozenset( 

113 (row["table_name"], row["column_name"]) 

114 for row in await database.query_raw( 

115 "SELECT table_name, column_name FROM information_schema.columns " 

116 "WHERE table_schema = ANY (current_schemas(false))" 

117 ) 

118 ) 

119 return tuple( 

120 secret_column for secret_column in _SECRET_COLUMNS if (secret_column.table, secret_column.column) in existing 

121 ) 

122 

123 

124async def _rows_of(database: SupportsRawQueries, secret_column: _SecretColumn) -> tuple[Mapping[str, object], ...]: 

125 marked_only: Final = ( 

126 f" AND \"{secret_column.column}\"::text LIKE '%{CALLBACK_VAR_ENCRYPTED_PREFIX}%'" 

127 if secret_column.only_rows_with_marked_ciphertexts 

128 else "" 

129 ) 

130 return tuple( 

131 await database.query_raw( 

132 f'SELECT "{secret_column.primary_key}", "{secret_column.column}" FROM "{secret_column.table}" ' 

133 f'WHERE "{secret_column.column}" IS NOT NULL{marked_only}' 

134 ) 

135 ) 

136 

137 

138async def _reencrypt_row( 

139 database: SupportsRawQueries, 

140 secret_column: _SecretColumn, 

141 row: Mapping[str, object], 

142 reencrypted: ReplaceCiphertext, 

143) -> int: 

144 stored: Final = _STORED_VALUE.validate_python(row[secret_column.column]) 

145 migrated, count = replace_ciphertexts(stored, reencrypted) 

146 if count == 0: 

147 return 0 

148 cast_to: Final = "::jsonb" if secret_column.is_json else "" 

149 rows_updated: Final = await database.execute_raw( 

150 f'UPDATE "{secret_column.table}" SET "{secret_column.column}" = $1{cast_to} ' 

151 f'WHERE "{secret_column.primary_key}" = $2 AND "{secret_column.column}" = $3{cast_to}', 

152 _as_sql_parameter(migrated, secret_column), 

153 _PRIMARY_KEY.validate_python(row[secret_column.primary_key]), 

154 _as_sql_parameter(stored, secret_column), 

155 ) 

156 return count if rows_updated else 0 

157 

158 

159def _as_sql_parameter(value: JsonValue, secret_column: _SecretColumn) -> str: 

160 return json.dumps(value) if secret_column.is_json else _CIPHERTEXT.validate_python(value) 

161 

162 

163class NothingToMigrate(Enum): 

164 SALT_KEY_ENCRYPTS_STORED_VALUES = "salt_key_encrypts_stored_values" 

165 NO_DATABASE = "no_database" 

166 NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY = "nothing_encrypted_with_previous_key" 

167 

168 

169@dataclass(frozen=True, slots=True) 

170class Migrated: 

171 migrated: int 

172 remaining: int 

173 

174 

175@dataclass(frozen=True, slots=True) 

176class MigrationFailed: 

177 error: Exception 

178 

179 

180MigrationOutcome = NothingToMigrate | Migrated | MigrationFailed 

181 

182 

183async def migrate_if_requested( 

184 *, 

185 environ: Mapping[str, str], 

186 master_key: str | None, 

187 connected_database: Callable[[], SupportsRawQueries | None], 

188 log: Callable[[str], None], 

189 raise_unless_tolerated: Callable[[Exception], None], 

190) -> MigrationOutcome | None: 

191 previous_master_key: Final = environ.get(MIGRATE_FROM_MASTER_KEY_ENV_VAR) 

192 if previous_master_key is None or master_key is None: 192 ↛ 194line 192 didn't jump to line 194 because the condition on line 192 was always true

193 return None 

194 outcome: Final = await migrate_from_previous_master_key( 

195 previous_master_key=previous_master_key, 

196 master_key=master_key, 

197 salt_key_is_set=SALT_KEY_ENV_VAR in environ, 

198 database=connected_database(), 

199 log=log, 

200 ) 

201 if isinstance(outcome, MigrationFailed): 

202 raise_unless_tolerated(outcome.error) 

203 return outcome 

204 

205 

206async def migrate_from_previous_master_key( 

207 *, 

208 previous_master_key: str, 

209 master_key: str, 

210 salt_key_is_set: bool, 

211 database: SupportsRawQueries | None, 

212 log: Callable[[str], None], 

213) -> MigrationOutcome: 

214 outcome: Final = await _migrate_or_failure( 

215 previous_master_key=previous_master_key, 

216 master_key=master_key, 

217 salt_key_is_set=salt_key_is_set, 

218 database=database, 

219 log=log, 

220 ) 

221 log(describe_outcome(outcome)) 

222 return outcome 

223 

224 

225async def _migrate_or_failure( 

226 *, 

227 previous_master_key: str, 

228 master_key: str, 

229 salt_key_is_set: bool, 

230 database: SupportsRawQueries | None, 

231 log: Callable[[str], None], 

232) -> MigrationOutcome: 

233 try: 

234 return await _migrate( 

235 previous_master_key=previous_master_key, 

236 master_key=master_key, 

237 salt_key_is_set=salt_key_is_set, 

238 database=database, 

239 log=log, 

240 ) 

241 except Exception as error: # noqa: BLE001 # a value, so the boot applies its own database outage rule to it 

242 return MigrationFailed(error=error) 

243 

244 

245async def _migrate( 

246 *, 

247 previous_master_key: str, 

248 master_key: str, 

249 salt_key_is_set: bool, 

250 database: SupportsRawQueries | None, 

251 log: Callable[[str], None], 

252) -> MigrationOutcome: 

253 if salt_key_is_set: 

254 return NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES 

255 if database is None: 

256 return NothingToMigrate.NO_DATABASE 

257 if previous_master_key == master_key: 

258 return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY 

259 found: Final = await count_values_encrypted_with(database, previous_master_key) 

260 if found == 0: 

261 return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY 

262 log(f"Re-encrypting {found} stored value(s) from the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} key to the new master key.") 

263 migrated: Final = Migrated( 

264 migrated=await reencrypt_stored_values(database, from_key=previous_master_key, to_key=master_key), 

265 remaining=await count_values_encrypted_with(database, previous_master_key), 

266 ) 

267 another_worker_migrated_everything: Final = migrated == Migrated(migrated=0, remaining=0) 

268 return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY if another_worker_migrated_everything else migrated 

269 

270 

271def describe_outcome(outcome: MigrationOutcome) -> str: 

272 match outcome: 

273 case NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES: 

274 return ( 

275 f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is set, but {SALT_KEY_ENV_VAR} is what encrypts your stored " 

276 f"values, so there is nothing to migrate. You may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}." 

277 ) 

278 case NothingToMigrate.NO_DATABASE: 

279 return ( 

280 f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is set, but no database is connected, so nothing was migrated. If " 

281 f"this proxy has no database, you may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}." 

282 ) 

283 case NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY: 

284 return ( 

285 f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is still set, but nothing in the database is left to migrate " 

286 f"from that key. You may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}." 

287 ) 

288 case Migrated(migrated=migrated, remaining=0): 

289 return ( 

290 f"Done re-encrypting {migrated} stored value(s) with the new master key. You may now delete the " 

291 f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} environment variable." 

292 ) 

293 case Migrated(migrated=migrated, remaining=remaining): 

294 return ( 

295 f"Re-encrypted {migrated} stored value(s), but {remaining} are still encrypted with the previous key " 

296 f"because they changed during the migration. Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set and restart " 

297 "the proxy to migrate them." 

298 ) 

299 case MigrationFailed(error=error): 

300 cause: Final = f"{type(error).__name__}: {error}"[:300] 

301 return ( 

302 f"Could not migrate stored values from the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} key ({cause}). Values " 

303 "still encrypted with the previous key cannot be read until the migration succeeds. Keep " 

304 f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} set and restart the proxy once the database is reachable." 

305 ) 

306 case _: 

307 assert_never(outcome)