mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(aws_secret_manager_v2): restore secret scheduled for deletion instead of failing CreateSecret (#42454)
* fix(aws_secret_manager_v2): restore secret scheduled for deletion instead of failing CreateSecret Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(aws_secret_manager_v2): reapply CreateSecret metadata and reschedule deletion when in-place restore fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
5035c458fb
commit
64456ce103
2 changed files with 286 additions and 23 deletions
|
|
@ -16,6 +16,8 @@ Requires:
|
|||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -294,32 +296,13 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
raise ValueError("Tags must be a dict or list of {Key, Value} pairs")
|
||||
data["Tags"] = tags_list
|
||||
|
||||
endpoint_url, headers, body = self._prepare_request(
|
||||
action="CreateSecret",
|
||||
create_response: Final = await self._async_create_or_restore_secret(
|
||||
secret_name=secret_name,
|
||||
secret_value=secret_value,
|
||||
optional_params=optional_params,
|
||||
request_data=data,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async_client: Final = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.SecretManager,
|
||||
params={"timeout": timeout},
|
||||
)
|
||||
|
||||
try:
|
||||
response: Final = await async_client.post(
|
||||
url=endpoint_url,
|
||||
headers=headers,
|
||||
data=body.decode("utf-8"),
|
||||
)
|
||||
response.raise_for_status()
|
||||
create_response: Final = response.json()
|
||||
except httpx.HTTPStatusError as err:
|
||||
raise ValueError(f"HTTP error occurred: {err.response.text}")
|
||||
except httpx.TimeoutException:
|
||||
raise ValueError("Timeout error occurred")
|
||||
|
||||
if self.replica_regions:
|
||||
try:
|
||||
await self.async_replicate_secret(
|
||||
|
|
@ -343,6 +326,110 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
|
||||
return create_response
|
||||
|
||||
async def _async_create_or_restore_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
request_data: Mapping[str, object],
|
||||
optional_params: dict | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
) -> dict[str, object]:
|
||||
try:
|
||||
return await self._async_post_action(
|
||||
action="CreateSecret",
|
||||
secret_name=secret_name,
|
||||
request_data=request_data,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
)
|
||||
except ValueError:
|
||||
if not await self._async_is_scheduled_for_deletion(
|
||||
secret_name=secret_name,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
):
|
||||
raise
|
||||
|
||||
verbose_logger.info(
|
||||
"Secret %s is scheduled for deletion, restoring and updating in place (RestoreSecret + UpdateSecret)",
|
||||
secret_name,
|
||||
)
|
||||
await self._async_post_action(
|
||||
action="RestoreSecret",
|
||||
secret_name=secret_name,
|
||||
request_data=None,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
)
|
||||
update_data: Final = MappingProxyType(
|
||||
{("SecretId" if key == "Name" else key): value for key, value in request_data.items() if key != "Tags"}
|
||||
)
|
||||
tags: Final = request_data.get("Tags")
|
||||
try:
|
||||
updated: Final = await self._async_post_action(
|
||||
action="UpdateSecret",
|
||||
secret_name=secret_name,
|
||||
request_data=update_data,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
)
|
||||
if tags is not None:
|
||||
await self._async_post_action(
|
||||
action="TagResource",
|
||||
secret_name=secret_name,
|
||||
request_data=MappingProxyType({"SecretId": secret_name, "Tags": tags}),
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
)
|
||||
except ValueError:
|
||||
await self.async_delete_secret(secret_name=secret_name, optional_params=optional_params, timeout=timeout)
|
||||
raise
|
||||
return updated
|
||||
|
||||
async def _async_is_scheduled_for_deletion(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: dict | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
) -> bool:
|
||||
try:
|
||||
described: Final = await self._async_post_action(
|
||||
action="DescribeSecret",
|
||||
secret_name=secret_name,
|
||||
request_data=None,
|
||||
optional_params=optional_params,
|
||||
timeout=timeout,
|
||||
)
|
||||
except ValueError:
|
||||
return False
|
||||
return described.get("DeletedDate") is not None
|
||||
|
||||
async def _async_post_action(
|
||||
self,
|
||||
action: str,
|
||||
secret_name: str,
|
||||
request_data: Mapping[str, object] | None,
|
||||
optional_params: dict | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
) -> dict[str, object]:
|
||||
endpoint_url, headers, body = self._prepare_request(
|
||||
action=action,
|
||||
secret_name=secret_name,
|
||||
optional_params=optional_params,
|
||||
request_data=dict(request_data) if request_data is not None else None,
|
||||
)
|
||||
async_client: Final = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.SecretManager,
|
||||
params={"timeout": timeout},
|
||||
)
|
||||
try:
|
||||
response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8"))
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except httpx.HTTPStatusError as err:
|
||||
raise ValueError(f"HTTP error occurred: {err.response.text}")
|
||||
except httpx.TimeoutException:
|
||||
raise ValueError("Timeout error occurred")
|
||||
|
||||
async def async_replicate_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
|
|
|
|||
|
|
@ -1,11 +1,17 @@
|
|||
from collections.abc import Mapping
|
||||
import json
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
|
||||
OptionalParams: TypeAlias = Mapping[str, object] | None
|
||||
|
|
@ -204,3 +210,173 @@ async def test_rotate_secret_different_names_persists_requested_value_and_delete
|
|||
assert manager.storage.values[new_name] == new_value
|
||||
assert current_name not in manager.storage.values
|
||||
assert manager.storage.values[unrelated_secret_name] == unrelated_value
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FakeSecretsManagerState:
|
||||
live: Mapping[str, str]
|
||||
scheduled_for_deletion: frozenset[str] = frozenset()
|
||||
actions: tuple[str, ...] = ()
|
||||
descriptions: Mapping[str, str] = MappingProxyType({})
|
||||
failing_actions: frozenset[str] = frozenset()
|
||||
|
||||
|
||||
class FakeSecretsManagerService:
|
||||
def __init__(self, state: FakeSecretsManagerState) -> None:
|
||||
self.state = state
|
||||
|
||||
def handle(self, request: httpx.Request) -> httpx.Response:
|
||||
action: Final = request.headers["X-Amz-Target"].removeprefix("secretsmanager.")
|
||||
body: Final = json.loads(request.content)
|
||||
name: Final = str(body.get("Name") or body.get("SecretId"))
|
||||
self.state = replace(self.state, actions=(*self.state.actions, f"{action}:{name}"))
|
||||
if action in self.state.failing_actions:
|
||||
return self._error("InternalServiceError", f"injected failure for {action}")
|
||||
match action:
|
||||
case "CreateSecret":
|
||||
if name in self.state.live:
|
||||
return self._error("ResourceExistsException", f"The secret {name} already exists")
|
||||
self.state = replace(
|
||||
self.state,
|
||||
live=MappingProxyType({**self.state.live, name: str(body["SecretString"])}),
|
||||
descriptions=MappingProxyType({**self.state.descriptions, name: str(body.get("Description", ""))}),
|
||||
)
|
||||
return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name})
|
||||
case "UpdateSecret":
|
||||
if name in self.state.scheduled_for_deletion:
|
||||
return self._error(
|
||||
"InvalidRequestException",
|
||||
"You can't perform this operation on the secret because it was marked for deletion.",
|
||||
)
|
||||
self.state = replace(
|
||||
self.state,
|
||||
live=MappingProxyType({**self.state.live, name: str(body["SecretString"])}),
|
||||
descriptions=MappingProxyType({**self.state.descriptions, name: str(body.get("Description", ""))}),
|
||||
)
|
||||
return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name})
|
||||
case "DescribeSecret":
|
||||
if name not in self.state.live:
|
||||
return self._error("ResourceNotFoundException", "Secrets Manager can't find the specified secret.")
|
||||
deleted: Final = "2026-01-01T00:00:00Z" if name in self.state.scheduled_for_deletion else None
|
||||
return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name, "DeletedDate": deleted})
|
||||
case "RestoreSecret":
|
||||
self.state = replace(self.state, scheduled_for_deletion=self.state.scheduled_for_deletion - {name})
|
||||
return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name})
|
||||
case "PutSecretValue":
|
||||
if name in self.state.scheduled_for_deletion:
|
||||
return self._error(
|
||||
"InvalidRequestException",
|
||||
"You can't perform this operation on the secret because it was marked for deletion.",
|
||||
)
|
||||
self.state = replace(
|
||||
self.state, live=MappingProxyType({**self.state.live, name: str(body["SecretString"])})
|
||||
)
|
||||
return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name})
|
||||
case "GetSecretValue":
|
||||
if name not in self.state.live or name in self.state.scheduled_for_deletion:
|
||||
return self._error("ResourceNotFoundException", "Secrets Manager can't find the specified secret.")
|
||||
return httpx.Response(200, json={"SecretString": self.state.live[name]})
|
||||
case "DeleteSecret":
|
||||
self.state = replace(self.state, scheduled_for_deletion=self.state.scheduled_for_deletion | {name})
|
||||
return httpx.Response(200, json={"ARN": f"arn:fake:{name}", "Name": name})
|
||||
return self._error("UnsupportedAction", action)
|
||||
|
||||
@staticmethod
|
||||
def _error(error_type: str, message: str) -> httpx.Response:
|
||||
return httpx.Response(400, json={"__type": error_type, "message": message})
|
||||
|
||||
|
||||
@contextmanager
|
||||
def fake_secrets_manager(monkeypatch: pytest.MonkeyPatch) -> Iterator[FakeSecretsManagerService]:
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "synthetic-access-key")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "synthetic-secret-key")
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
|
||||
service: Final = FakeSecretsManagerService(FakeSecretsManagerState(live=MappingProxyType({})))
|
||||
cache_key: Final = "async_httpx_clienttimeout_None" + httpxSpecialProvider.SecretManager
|
||||
litellm.in_memory_llm_clients_cache.set_cache(
|
||||
key=cache_key,
|
||||
value=AsyncHTTPHandler(transport=httpx.MockTransport(service.handle)),
|
||||
)
|
||||
try:
|
||||
yield service
|
||||
finally:
|
||||
litellm.in_memory_llm_clients_cache.delete_cache(
|
||||
litellm.in_memory_llm_clients_cache.update_cache_key_with_event_loop(cache_key)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rotate_secret_back_to_name_inside_recovery_window_restores_and_stores_new_value(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
alias_a: Final = "synthetic/alias-a"
|
||||
alias_b: Final = "synthetic/alias-b"
|
||||
with fake_secrets_manager(monkeypatch) as fake:
|
||||
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
|
||||
await manager.async_write_secret(secret_name=alias_a, secret_value="value-1")
|
||||
await manager.async_rotate_secret(
|
||||
current_secret_name=alias_a, new_secret_name=alias_b, new_secret_value="value-2"
|
||||
)
|
||||
assert alias_a in fake.state.scheduled_for_deletion
|
||||
|
||||
await manager.async_rotate_secret(
|
||||
current_secret_name=alias_b, new_secret_name=alias_a, new_secret_value="value-3"
|
||||
)
|
||||
|
||||
assert await manager.async_read_secret(secret_name=alias_a) == "value-3"
|
||||
assert await manager.async_read_secret(secret_name=alias_b) is None
|
||||
assert fake.state.scheduled_for_deletion == frozenset({alias_b})
|
||||
assert fake.state.descriptions[alias_a] == f"Rotated from {alias_b}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_secret_to_name_inside_recovery_window_reschedules_deletion_when_update_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
alias: Final = "synthetic/deleted-alias"
|
||||
with fake_secrets_manager(monkeypatch) as fake:
|
||||
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
|
||||
await manager.async_write_secret(secret_name=alias, secret_value="value-1")
|
||||
await manager.async_delete_secret(secret_name=alias, recovery_window_in_days=7)
|
||||
fake.state = replace(fake.state, failing_actions=frozenset({"UpdateSecret"}))
|
||||
|
||||
with pytest.raises(ValueError, match="injected failure for UpdateSecret"):
|
||||
await manager.async_write_secret(secret_name=alias, secret_value="value-2")
|
||||
|
||||
assert fake.state.scheduled_for_deletion == frozenset({alias})
|
||||
assert fake.state.live[alias] == "value-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_secret_to_name_inside_recovery_window_restores_and_stores_new_value(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
alias: Final = "synthetic/deleted-alias"
|
||||
with fake_secrets_manager(monkeypatch) as fake:
|
||||
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
|
||||
await manager.async_write_secret(secret_name=alias, secret_value="value-1")
|
||||
await manager.async_delete_secret(secret_name=alias, recovery_window_in_days=7)
|
||||
|
||||
assert await manager.async_write_secret(secret_name=alias, secret_value="value-2") == {
|
||||
"ARN": f"arn:fake:{alias}",
|
||||
"Name": alias,
|
||||
}
|
||||
|
||||
assert await manager.async_read_secret(secret_name=alias) == "value-2"
|
||||
assert fake.state.scheduled_for_deletion == frozenset()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_secret_to_live_existing_name_still_fails_without_overwriting(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
alias: Final = "synthetic/live-alias"
|
||||
with fake_secrets_manager(monkeypatch) as fake:
|
||||
manager: Final = AWSSecretsManagerV2(aws_region_name="us-east-1")
|
||||
await manager.async_write_secret(secret_name=alias, secret_value="value-1")
|
||||
|
||||
with pytest.raises(ValueError, match="ResourceExistsException"):
|
||||
await manager.async_write_secret(secret_name=alias, secret_value="value-2")
|
||||
|
||||
assert await manager.async_read_secret(secret_name=alias) == "value-1"
|
||||
assert f"RestoreSecret:{alias}" not in fake.state.actions
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue