update key name when regenerating a key

This commit is contained in:
Ishaan Jaff 2024-08-26 18:40:51 -07:00
parent 3527e47b59
commit 4dc2eea58d
2 changed files with 8 additions and 1 deletions

View file

@ -1027,10 +1027,15 @@ async def regenerate_key_fn(
new_token = f"sk-{secrets.token_urlsafe(16)}" new_token = f"sk-{secrets.token_urlsafe(16)}"
new_token_hash = hash_token(new_token) new_token_hash = hash_token(new_token)
new_token_key_name = f"sk-...{new_token[-4:]}"
# update new token in DB # update new token in DB
updated_token = await prisma_client.db.litellm_verificationtoken.update( updated_token = await prisma_client.db.litellm_verificationtoken.update(
where={"token": hashed_api_key}, data={"token": new_token_hash} where={"token": hashed_api_key},
data={
"token": new_token_hash,
"key_name": new_token_key_name,
},
) )
updated_token_dict = {} updated_token_dict = {}
if updated_token is not None: if updated_token is not None:

View file

@ -3035,4 +3035,6 @@ async def test_regenerate_api_key(prisma_client):
new_key.models == models new_key.models == models
), f"Expected models {models} but got {new_key.models}" ), f"Expected models {models} but got {new_key.models}"
assert new_key.key_name == f"sk-...{new_key.key[-4:]}"
pass pass