mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
fix(e2e): delete raw cloud-storage batch files with the master key
DELETE /v1/files/{id} only lets a proxy admin key delete a raw s3:// or
gs:// file id, because such ids skip the managed-file owner check. The
batch lifecycle cleanup deleted the vertex_ai raw ids with the test's
own virtual key and got a 403 at teardown on every build since #194
Raw cloud-storage ids now go through the master key; managed and
provider-native ids keep using the creating key
This commit is contained in:
parent
1d71564063
commit
dd6ef9e1bc
4 changed files with 28 additions and 2 deletions
|
|
@ -5,7 +5,7 @@ from time import monotonic, sleep
|
|||
from typing import Final, Protocol
|
||||
|
||||
from batch_client import BatchObject, FileDeleteResponse
|
||||
from capabilities import is_managed_id
|
||||
from capabilities import is_cloud_storage_id, is_managed_id
|
||||
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -19,6 +19,8 @@ BATCH_CANCEL_POLL_SECONDS: Final = 10.0
|
|||
class BatchCleanupClient(Protocol):
|
||||
def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: ...
|
||||
|
||||
def delete_file_as_admin(self, file_id: str, *, provider: str | None = None) -> Result[FileDeleteResponse]: ...
|
||||
|
||||
def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...
|
||||
|
||||
def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...
|
||||
|
|
@ -49,7 +51,12 @@ def _require_cleanup_success[R: BaseModel](result: Result[R], operation: str) ->
|
|||
|
||||
|
||||
def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None = None) -> None:
|
||||
result: Final = cleanup_result(lambda: client.delete_file(file_id, key=key, provider=provider))
|
||||
delete: Final[Callable[[], Result[FileDeleteResponse]]] = (
|
||||
(lambda: client.delete_file_as_admin(file_id, provider=provider))
|
||||
if is_cloud_storage_id(file_id)
|
||||
else (lambda: client.delete_file(file_id, key=key, provider=provider))
|
||||
)
|
||||
result: Final = cleanup_result(delete)
|
||||
if isinstance(result, UnknownApiError) and result.status_code == 404:
|
||||
return
|
||||
deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}")
|
||||
|
|
|
|||
|
|
@ -233,6 +233,14 @@ class BatchClient:
|
|||
response_type=FileDeleteResponse,
|
||||
)
|
||||
|
||||
def delete_file_as_admin(self, file_id: str, *, provider: str | None = None) -> Result[FileDeleteResponse]:
|
||||
return self.proxy.transport.delete(
|
||||
f"{_files_path(provider)}/{file_id}",
|
||||
headers=self.proxy.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=FileDeleteResponse,
|
||||
)
|
||||
|
||||
|
||||
def _files_path(provider: str | None) -> str:
|
||||
return f"/{provider}/v1/files" if provider else "/v1/files"
|
||||
|
|
|
|||
|
|
@ -222,6 +222,13 @@ def is_managed_id(id_str: str) -> bool:
|
|||
return _b64_decode(id_str).startswith("litellm_proxy")
|
||||
|
||||
|
||||
CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://")
|
||||
|
||||
|
||||
def is_cloud_storage_id(id_str: str) -> bool:
|
||||
return id_str.startswith(CLOUD_STORAGE_SCHEMES)
|
||||
|
||||
|
||||
def is_model_encoded_id(id_str: str) -> bool:
|
||||
for prefix in ("file-", "batch_"):
|
||||
if id_str.startswith(prefix):
|
||||
|
|
|
|||
|
|
@ -45,6 +45,10 @@ class CleanupClient:
|
|||
self.calls(f"delete {provider} {file_id}")
|
||||
return self.file_response()
|
||||
|
||||
def delete_file_as_admin(self, file_id: str, *, provider: str | None = None) -> Result[FileDeleteResponse]:
|
||||
self.calls(f"admin delete {provider} {file_id}")
|
||||
return self.file_response()
|
||||
|
||||
def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
|
||||
self.calls(f"retrieve {provider} {batch_id}")
|
||||
return self.batch_response()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue