mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(proxy-extras): log v1 migration failures at ERROR so LITELLM_LOG=ERROR shows them (#44202)
* fix(proxy-extras): log v1 migration failures at ERROR so LITELLM_LOG=ERROR shows them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy-extras): reuse litellm secret redaction and mask configured DB passwords exactly Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy-extras): name the password alternation in _redact_credentials Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy-extras): wrap v1 migration ERROR lines at 120 columns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy-extras): integration cells for v1 migration ERROR logging and password redaction Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy-extras): integration cells for component DB env vars, JSON logs, migration Job and v2 resolver Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy-extras): give every v1 migration integration cell the 900s timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy-extras): bind the recovery forwarder before migrating so the retry cannot race it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy-extras): drop the slow P3005 integration cell and bound migration subprocesses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy-extras): keep command repr and tolerate non-sequence cmd in migration error logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy-extras): double the subprocess boundary in the cmd=None retry test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: jesus <jesus@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng <yucheng@berri.ai>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
cf22deb96a
commit
984b4134be
3 changed files with 1048 additions and 5 deletions
|
|
@ -1,3 +1,4 @@
|
|||
import functools
|
||||
import glob
|
||||
import os
|
||||
import random
|
||||
|
|
@ -10,7 +11,8 @@ import time
|
|||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
from typing import TYPE_CHECKING, Final, Optional, Union
|
||||
from urllib.parse import unquote, urlsplit
|
||||
|
||||
from litellm_proxy_extras import prisma_toolchain
|
||||
from litellm_proxy_extras._logging import logger
|
||||
|
|
@ -202,6 +204,66 @@ def _max_migration_timestamp(names) -> int:
|
|||
return max(_migration_timestamp(n) for n in names)
|
||||
|
||||
|
||||
_REDACTED: Final = "REDACTED"
|
||||
_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _secret_shape_redactor() -> Callable[[str], str]:
|
||||
try:
|
||||
from litellm._logging import redact_secrets
|
||||
except ImportError:
|
||||
return lambda text: text
|
||||
return redact_secrets
|
||||
|
||||
|
||||
def _url_passwords(url: str) -> frozenset[str]:
|
||||
try:
|
||||
parts: Final = urlsplit(url)
|
||||
except ValueError:
|
||||
return frozenset()
|
||||
query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
|
||||
raw_query_passwords: Final = tuple(
|
||||
value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
|
||||
)
|
||||
raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
|
||||
return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
|
||||
|
||||
|
||||
def _configured_database_passwords() -> frozenset[str]:
|
||||
database_url: Final = os.getenv("DATABASE_URL")
|
||||
direct_url: Final = os.getenv("DIRECT_URL")
|
||||
database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
|
||||
direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
|
||||
return database_passwords | direct_passwords
|
||||
|
||||
|
||||
def _redact_credentials(text: str) -> str:
|
||||
"""Mask configured database passwords before passing the text to LiteLLM redaction."""
|
||||
passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
|
||||
alternation: Final = "|".join(re.escape(password) for password in passwords)
|
||||
password_pattern: Final = (
|
||||
re.compile(rf"(?P<lead>:|password=)(?:{alternation})(?=@|&|$|[\s'\"\]),])", re.IGNORECASE)
|
||||
if passwords
|
||||
else None
|
||||
)
|
||||
result: Final = password_pattern.sub(rf"\g<lead>{_REDACTED}", text) if password_pattern is not None else text
|
||||
return _secret_shape_redactor()(result)
|
||||
|
||||
|
||||
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
|
||||
if isinstance(command, tuple):
|
||||
return tuple(_redact_credentials(str(argument)) for argument in command)
|
||||
if isinstance(command, list):
|
||||
return [_redact_credentials(str(argument)) for argument in command]
|
||||
return _redact_credentials(str(command))
|
||||
|
||||
|
||||
def _redact_command_error(error: subprocess.CalledProcessError) -> str:
|
||||
redacted_command: Final = _redacted_command(error.cmd)
|
||||
return str(subprocess.CalledProcessError(error.returncode, redacted_command))
|
||||
|
||||
|
||||
def _get_prisma_command() -> str:
|
||||
"""Get the Prisma command to use, bypassing Python wrapper in offline mode."""
|
||||
if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
|
||||
|
|
@ -315,7 +377,8 @@ class ProxyExtrasDBManager:
|
|||
return False
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(
|
||||
f"Error creating baseline migration: {e}, {e.stderr}, {e.stdout}"
|
||||
f"Error creating baseline migration: {_redact_command_error(e)}, "
|
||||
f"{_redact_credentials(str(e.stderr))}, {_redact_credentials(str(e.stdout))}"
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1572,6 +1635,11 @@ class ProxyExtrasDBManager:
|
|||
f"Error: {stderr}"
|
||||
)
|
||||
raise
|
||||
else:
|
||||
logger.error(
|
||||
"prisma migrate deploy failed with an error the resolver does not handle: "
|
||||
f"{_redact_credentials(stderr)}"
|
||||
)
|
||||
else:
|
||||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
|
||||
|
|
@ -1586,7 +1654,7 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
return True
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(
|
||||
logger.error(
|
||||
"Attempt %s timed out. Raise %s if this database needs longer to apply its schema.",
|
||||
attempt + 1,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
|
|
@ -1599,7 +1667,12 @@ class ProxyExtrasDBManager:
|
|||
if attempts_left > 0
|
||||
else ""
|
||||
)
|
||||
logger.info(f"The process failed to execute. Details: {e}.{retry_msg}")
|
||||
stderr_detail: Final = (
|
||||
f" stderr: {_redact_credentials(str(e.stderr))}" if e.stderr else ""
|
||||
)
|
||||
logger.error(
|
||||
f"The process failed to execute. Details: {_redact_command_error(e)}.{stderr_detail}{retry_msg}"
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
|
|||
558
tests/integration/database/test_v1_migration_error_logging.py
Normal file
558
tests/integration/database/test_v1_migration_error_logging.py
Normal file
|
|
@ -0,0 +1,558 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import socket
|
||||
import socketserver
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from contextlib import contextmanager, suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, cast
|
||||
from urllib.parse import quote, urlsplit, urlunsplit
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from psycopg import sql
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
from tests.integration._support.client import eventually
|
||||
from tests.integration._support.process import _free_port as free_port
|
||||
|
||||
REPO_ROOT: Final = Path(__file__).resolve().parents[3]
|
||||
MIGRATIONS_DIR: Final = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
|
||||
MIGRATION_NAME: Final = "20260921190000_agent_identity"
|
||||
pytestmark: Final = pytest.mark.timeout(300)
|
||||
PASSWORD: Final = "wr ong'pw9"
|
||||
FRAGMENTS: Final = ("wr ong", "wr+ong", "ong'pw9", "wr%20ong", "ong%27pw9", "pw9")
|
||||
SHIPPED_MIGRATIONS: Final = tuple(
|
||||
sorted(path.name for path in MIGRATIONS_DIR.iterdir() if path.is_dir() and path.name != "0_init")
|
||||
)
|
||||
LOG_PREFIX: Final = r"^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2},\d{3} - [^\n]+ - "
|
||||
LOG_RECORD_START: Final = rf"{LOG_PREFIX}(?:DEBUG|INFO|WARNING|ERROR|CRITICAL) - "
|
||||
ERROR_RECORD: Final = re.compile(rf"(?ms)^({LOG_PREFIX}ERROR - .*?)(?={LOG_RECORD_START}|\Z)")
|
||||
RETRY_COUNT: Final = re.compile(r"Retrying\.\.\. \((\d+) attempts left\)")
|
||||
FRAGMENT_PATTERN: Final = re.compile(
|
||||
"|".join(re.escape(fragment) for fragment in sorted(FRAGMENTS, key=len, reverse=True))
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MigrationResult:
|
||||
returncode: int
|
||||
output: str
|
||||
|
||||
|
||||
def _migration_log_path(test_name: str, tmp_path: Path) -> Path:
|
||||
log_directory: Final = (
|
||||
Path(os.environ["INTEGRATION_RESULTS_DIR"]) if "INTEGRATION_RESULTS_DIR" in os.environ else tmp_path
|
||||
)
|
||||
log_path: Final = log_directory / f"v1-migration-{test_name}-{uuid.uuid4().hex[:8]}.log"
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
return log_path
|
||||
|
||||
|
||||
def _database_url(
|
||||
admin_url: str,
|
||||
role: str,
|
||||
password: str,
|
||||
database: str,
|
||||
*,
|
||||
encode_password: bool = True,
|
||||
) -> str:
|
||||
parsed: Final = urlsplit(admin_url)
|
||||
authority: Final = parsed.netloc.rsplit("@", 1)[-1]
|
||||
encoded_password: Final = quote(password, safe="") if encode_password else password
|
||||
netloc: Final = f"{quote(role, safe='')}:{encoded_password}@{authority}"
|
||||
return urlunsplit(parsed._replace(netloc=netloc, path=f"/{database}"))
|
||||
|
||||
|
||||
def _replace_port(database_url: str, port: int, hostname: str | None = None) -> str:
|
||||
parsed: Final = urlsplit(database_url)
|
||||
target_host: Final = hostname or parsed.hostname
|
||||
assert target_host is not None
|
||||
host: Final = f"[{target_host}]" if ":" in target_host else target_host
|
||||
userinfo: Final = parsed.netloc.rsplit("@", 1)[0]
|
||||
return urlunsplit(parsed._replace(netloc=f"{userinfo}@{host}:{port}"))
|
||||
|
||||
|
||||
def _unreachable_database_url(database_url: str) -> str:
|
||||
parsed: Final = urlsplit(database_url)
|
||||
url: Final = _database_url(
|
||||
database_url,
|
||||
parsed.username or "",
|
||||
PASSWORD,
|
||||
parsed.path.lstrip("/"),
|
||||
encode_password=False,
|
||||
)
|
||||
return _replace_port(url, free_port())
|
||||
|
||||
|
||||
@contextmanager
|
||||
def owned_database(password: str) -> Iterator[str]:
|
||||
admin_url: Final = os.environ["DATABASE_URL"]
|
||||
role: Final = f"v1_migration_{uuid.uuid4().hex}"
|
||||
database: Final = f"v1_migration_{uuid.uuid4().hex}"
|
||||
database_url: Final = _database_url(admin_url, role, password, database)
|
||||
try:
|
||||
with psycopg.connect(admin_url, autocommit=True) as admin:
|
||||
admin.execute(
|
||||
sql.SQL("CREATE ROLE {} WITH LOGIN PASSWORD {}").format(sql.Identifier(role), sql.Literal(password))
|
||||
)
|
||||
admin.execute(sql.SQL("CREATE DATABASE {} OWNER {}").format(sql.Identifier(database), sql.Identifier(role)))
|
||||
yield database_url
|
||||
finally:
|
||||
with psycopg.connect(admin_url, autocommit=True) as admin:
|
||||
admin.execute(sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(sql.Identifier(database)))
|
||||
admin.execute(sql.SQL("DROP ROLE IF EXISTS {}").format(sql.Identifier(role)))
|
||||
|
||||
|
||||
def _migration_environment(database_url: str | None, extra_env: Mapping[str, str]) -> dict[str, str]:
|
||||
excluded_variables: Final = (
|
||||
("DIRECT_URL", "USE_V2_MIGRATION_RESOLVER")
|
||||
if database_url is not None
|
||||
else ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER")
|
||||
)
|
||||
inherited_environment: Final = {key: value for key, value in os.environ.items() if key not in excluded_variables}
|
||||
database_environment: Final = {"DATABASE_URL": database_url} if database_url is not None else {}
|
||||
return {
|
||||
**inherited_environment,
|
||||
**database_environment,
|
||||
"LITELLM_LOG": "ERROR",
|
||||
**extra_env,
|
||||
}
|
||||
|
||||
|
||||
def _migration_invocation(
|
||||
database_url: str | None,
|
||||
tmp_path: Path,
|
||||
extra_env: Mapping[str, str],
|
||||
resolver: Literal["legacy", "v2"] = "legacy",
|
||||
) -> tuple[tuple[str, ...], dict[str, str]]:
|
||||
config_path: Final = tmp_path / "config.yaml"
|
||||
config_path.write_text(
|
||||
"model_list:\n - model_name: integration-fake\n litellm_params:\n model: openai/integration-fake\n"
|
||||
)
|
||||
resolver_flag: Final = "--use_legacy_migration_resolver" if resolver == "legacy" else "--use_v2_migration_resolver"
|
||||
command: Final = (
|
||||
sys.executable,
|
||||
"-I",
|
||||
"-m",
|
||||
"litellm.proxy.proxy_cli",
|
||||
"--config",
|
||||
str(config_path),
|
||||
resolver_flag,
|
||||
"--skip_server_startup",
|
||||
)
|
||||
environment: Final = _migration_environment(database_url, extra_env)
|
||||
return command, environment
|
||||
|
||||
|
||||
def run_v1_migrations(
|
||||
database_url: str | None,
|
||||
tmp_path: Path,
|
||||
extra_env: Mapping[str, str],
|
||||
test_name: str,
|
||||
resolver: Literal["legacy", "v2"] = "legacy",
|
||||
) -> MigrationResult:
|
||||
command, environment = _migration_invocation(database_url, tmp_path, extra_env, resolver)
|
||||
output_path: Final = _migration_log_path(test_name, tmp_path)
|
||||
with output_path.open("w") as output_file:
|
||||
completed: Final = subprocess.run(
|
||||
command,
|
||||
cwd=REPO_ROOT,
|
||||
env=environment,
|
||||
stdout=output_file,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
timeout=240,
|
||||
)
|
||||
output: Final = output_path.read_text()
|
||||
return MigrationResult(completed.returncode, output)
|
||||
|
||||
|
||||
def error_lines(output: str) -> tuple[str, ...]:
|
||||
return tuple(match.group(1) for match in ERROR_RECORD.finditer(output))
|
||||
|
||||
|
||||
def _json_error_records(output: str) -> tuple[dict[str, object], ...]:
|
||||
records: Final = tuple(_parse_json_record(line, output) for line in output.splitlines() if line.startswith("{"))
|
||||
return tuple(record for record in records if record.get("level") == "ERROR")
|
||||
|
||||
|
||||
def _parse_json_record(line: str, output: str) -> dict[str, object]:
|
||||
try:
|
||||
record: Final = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
pytest.fail(_safe_output(output))
|
||||
assert isinstance(record, dict), _safe_output(output)
|
||||
return cast(dict[str, object], record)
|
||||
|
||||
|
||||
def _retry_count_texts(lines: tuple[str, ...]) -> tuple[str, ...]:
|
||||
return tuple(value for value in (_retry_count_text(line) for line in lines) if value is not None)
|
||||
|
||||
|
||||
def _retry_count_text(line: str) -> str | None:
|
||||
match: Final = RETRY_COUNT.search(line)
|
||||
return match.group(1) if match is not None else None
|
||||
|
||||
|
||||
def _safe_output(output: str) -> str:
|
||||
return FRAGMENT_PATTERN.sub("[REDACTED]", output)
|
||||
|
||||
|
||||
def _assert_no_password_fragments(output: str) -> None:
|
||||
assert [fragment for fragment in FRAGMENTS if fragment in output] == [], _safe_output(output)
|
||||
|
||||
|
||||
def _applied_migrations(database_url: str) -> tuple[str, ...]:
|
||||
with psycopg.connect(database_url) as connection:
|
||||
rows: Final = connection.execute(
|
||||
'SELECT migration_name FROM "_prisma_migrations" '
|
||||
"WHERE finished_at IS NOT NULL AND rolled_back_at IS NULL ORDER BY migration_name"
|
||||
).fetchall()
|
||||
return tuple(str(row[0]) for row in rows)
|
||||
|
||||
|
||||
def _duplicate_migrations(database_url: str) -> tuple[tuple[str, int], ...]:
|
||||
with psycopg.connect(database_url) as connection:
|
||||
rows: Final = connection.execute(
|
||||
'SELECT migration_name, COUNT(*) FROM "_prisma_migrations" '
|
||||
"GROUP BY migration_name HAVING COUNT(*) > 1 ORDER BY migration_name"
|
||||
).fetchall()
|
||||
return tuple((str(row[0]), int(row[1])) for row in rows)
|
||||
|
||||
|
||||
def _retry_p3018_migration(database_url: str) -> None:
|
||||
agent_ids: Final = (uuid.uuid4().hex, uuid.uuid4().hex)
|
||||
with psycopg.connect(database_url) as connection:
|
||||
connection.execute('DROP INDEX "LiteLLM_AgentIdentity_provider_tenant_id_client_id_key"')
|
||||
for agent_id in agent_ids:
|
||||
connection.execute(
|
||||
'INSERT INTO "LiteLLM_AgentsTable" '
|
||||
'("agent_id", "agent_name", "agent_card_params", "created_by", "updated_by") '
|
||||
"VALUES (%s, %s, %s, %s, %s)",
|
||||
(agent_id, f"audit-agent-{agent_id}", Jsonb({}), "integration", "integration"),
|
||||
)
|
||||
connection.execute(
|
||||
'INSERT INTO "LiteLLM_AgentIdentity" '
|
||||
'("agent_id", "provider", "issuer", "tenant_id", "client_id", "revision") '
|
||||
"VALUES (%s, %s, %s, %s, %s, %s)",
|
||||
(agent_id, "entra", f"https://audit.invalid/{agent_id}", "tenant-1", "client-1", uuid.uuid4().hex),
|
||||
)
|
||||
connection.execute('DELETE FROM "_prisma_migrations" WHERE migration_name = %s', (MIGRATION_NAME,))
|
||||
|
||||
|
||||
def _relay(source: socket.socket, destination: socket.socket) -> None:
|
||||
try:
|
||||
while data := source.recv(65536):
|
||||
destination.sendall(data)
|
||||
except (BrokenPipeError, ConnectionAbortedError, ConnectionResetError):
|
||||
return
|
||||
|
||||
|
||||
class _PostgresForwardingServer(socketserver.ThreadingTCPServer):
|
||||
allow_reuse_address = True
|
||||
daemon_threads = True
|
||||
request_queue_size = 64
|
||||
target: tuple[str, int]
|
||||
|
||||
def __init__(self, port: int, target: tuple[str, int]) -> None:
|
||||
self.target = target
|
||||
super().__init__(("127.0.0.1", port), _PostgresForwardingHandler)
|
||||
|
||||
|
||||
class _PostgresForwardingHandler(socketserver.BaseRequestHandler):
|
||||
def handle(self) -> None:
|
||||
server: Final = cast(_PostgresForwardingServer, self.server)
|
||||
with socket.create_connection(server.target, timeout=10) as upstream:
|
||||
reply: Final = threading.Thread(target=_relay, args=(upstream, self.request), daemon=True)
|
||||
reply.start()
|
||||
try:
|
||||
_relay(self.request, upstream)
|
||||
finally:
|
||||
with suppress(OSError):
|
||||
self.request.shutdown(socket.SHUT_WR)
|
||||
with suppress(OSError):
|
||||
upstream.shutdown(socket.SHUT_WR)
|
||||
reply.join(timeout=10)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _gated_postgres_forwarder(port: int, target: tuple[str, int]) -> Iterator[Callable[[], None]]:
|
||||
server: Final = _PostgresForwardingServer(port, target)
|
||||
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
try:
|
||||
yield thread.start
|
||||
finally:
|
||||
if thread.is_alive():
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
if thread.ident is not None:
|
||||
thread.join(timeout=10)
|
||||
|
||||
|
||||
def _stop_process(process: subprocess.Popen[str]) -> None:
|
||||
if process.poll() is None:
|
||||
with suppress(ProcessLookupError):
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
try:
|
||||
process.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
with suppress(ProcessLookupError):
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
process.wait(timeout=30)
|
||||
|
||||
|
||||
def _migration_log_has_p1001_or_process_exited(output_path: Path, process: subprocess.Popen[str]) -> tuple[bool, bool]:
|
||||
p1001_logged: Final = any("P1001" in line for line in error_lines(output_path.read_text()))
|
||||
process_exited: Final = process.poll() is not None
|
||||
return p1001_logged, process_exited
|
||||
|
||||
|
||||
def test_unreachable_database_emits_four_p1001_errors_without_password_fragments(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
unreachable_url: Final = _unreachable_database_url(database_url)
|
||||
completed: Final = run_v1_migrations(unreachable_url, tmp_path, {}, request.node.name)
|
||||
output: Final = completed.output
|
||||
errors: Final = error_lines(output)
|
||||
p1001_errors: Final = tuple(line for line in errors if "P1001" in line)
|
||||
assert completed.returncode == 1, _safe_output(output)
|
||||
assert len(errors) == 4, _safe_output(output)
|
||||
assert len(p1001_errors) == 4, _safe_output(output)
|
||||
assert _retry_count_texts(errors) == (), _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_wrong_password_emits_four_p1000_errors_without_password_fragments(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(f"correct-{uuid.uuid4().hex}") as database_url:
|
||||
parsed: Final = urlsplit(database_url)
|
||||
wrong_url: Final = _database_url(
|
||||
database_url,
|
||||
parsed.username or "",
|
||||
PASSWORD,
|
||||
parsed.path.lstrip("/"),
|
||||
)
|
||||
completed: Final = run_v1_migrations(wrong_url, tmp_path, {}, request.node.name)
|
||||
output: Final = completed.output
|
||||
errors: Final = error_lines(output)
|
||||
p1000_errors: Final = tuple(line for line in errors if "P1000" in line)
|
||||
assert completed.returncode == 1, _safe_output(output)
|
||||
assert len(errors) == 4, _safe_output(output)
|
||||
assert len(p1000_errors) == 4, _safe_output(output)
|
||||
assert _retry_count_texts(errors) == (), _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_duplicate_agent_identity_logs_the_p3018_migration_error(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
setup: Final = run_v1_migrations(database_url, tmp_path, {}, request.node.name)
|
||||
setup_output: Final = setup.output
|
||||
assert setup.returncode == 0, _safe_output(setup_output)
|
||||
_retry_p3018_migration(database_url)
|
||||
completed: Final = run_v1_migrations(database_url, tmp_path, {}, request.node.name)
|
||||
output: Final = completed.output
|
||||
errors: Final = error_lines(output)
|
||||
p3018_errors: Final = tuple(line for line in errors if "P3018" in line)
|
||||
expected_markers: Final = ((True, True), (True, True))
|
||||
assert completed.returncode == 1, _safe_output(output)
|
||||
assert len(errors) == 2, _safe_output(output)
|
||||
assert tuple((MIGRATION_NAME in line, "P3018" in line) for line in p3018_errors) == expected_markers, (
|
||||
_safe_output(output)
|
||||
)
|
||||
assert _retry_count_texts(errors) == ("3", "1"), _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_clean_database_applies_exactly_the_shipped_migrations(tmp_path: Path, request: pytest.FixtureRequest) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
completed: Final = run_v1_migrations(database_url, tmp_path, {}, request.node.name)
|
||||
output: Final = completed.output
|
||||
assert completed.returncode == 0, _safe_output(output)
|
||||
assert error_lines(output) == (), _safe_output(output)
|
||||
assert _applied_migrations(database_url) == SHIPPED_MIGRATIONS, _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_unreachable_database_recovers_after_postgres_forwarder_starts(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
admin_url: Final = os.environ["DATABASE_URL"]
|
||||
admin: Final = urlsplit(admin_url)
|
||||
hostname: Final = admin.hostname
|
||||
port: Final = admin.port
|
||||
assert hostname is not None and port is not None
|
||||
target_host: Final = "127.0.0.1" if hostname == "localhost" else hostname
|
||||
forwarding_port: Final = free_port()
|
||||
forwarded_url: Final = _replace_port(database_url, forwarding_port, "127.0.0.1")
|
||||
command, environment = _migration_invocation(forwarded_url, tmp_path, {})
|
||||
output_path: Final = _migration_log_path(request.node.name, tmp_path)
|
||||
with _gated_postgres_forwarder(forwarding_port, (target_host, port)) as open_forwarder:
|
||||
with output_path.open("w") as output_file:
|
||||
process: Final = subprocess.Popen(
|
||||
command,
|
||||
cwd=REPO_ROOT,
|
||||
env=environment,
|
||||
stdout=output_file,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
start_new_session=True,
|
||||
)
|
||||
try:
|
||||
observation: Final = eventually(
|
||||
lambda: _migration_log_has_p1001_or_process_exited(output_path, process),
|
||||
lambda state: state[0] or state[1],
|
||||
seconds=240,
|
||||
)
|
||||
assert observation[0], _safe_output(output_path.read_text())
|
||||
open_forwarder()
|
||||
completed_returncode: Final = process.wait(timeout=240)
|
||||
finally:
|
||||
_stop_process(process)
|
||||
output: Final = output_path.read_text()
|
||||
errors: Final = error_lines(output)
|
||||
p1001_errors: Final = tuple(line for line in errors if "P1001" in line)
|
||||
assert completed_returncode == 0, _safe_output(output)
|
||||
assert len(p1001_errors) == 1, _safe_output(output)
|
||||
assert _retry_count_texts(errors) == (), _safe_output(output)
|
||||
assert _duplicate_migrations(database_url) == (), _safe_output(output)
|
||||
assert _applied_migrations(database_url) == SHIPPED_MIGRATIONS, _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_unreachable_database_keeps_password_masked_when_shape_redaction_is_disabled(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
unreachable_url: Final = _unreachable_database_url(database_url)
|
||||
completed: Final = run_v1_migrations(
|
||||
unreachable_url,
|
||||
tmp_path,
|
||||
{"LITELLM_DISABLE_REDACT_SECRETS": "true"},
|
||||
request.node.name,
|
||||
)
|
||||
output: Final = completed.output
|
||||
errors: Final = error_lines(output)
|
||||
p1001_errors: Final = tuple(line for line in errors if "P1001" in line)
|
||||
assert completed.returncode == 1, _safe_output(output)
|
||||
assert len(errors) == 4, _safe_output(output)
|
||||
assert len(p1001_errors) == 4, _safe_output(output)
|
||||
assert _retry_count_texts(errors) == (), _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_component_database_env_vars_with_wrong_password_emit_four_p1000_errors_without_password_fragments(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
correct_password: Final = f"correct-{uuid.uuid4().hex}"
|
||||
with owned_database(correct_password) as database_url:
|
||||
parsed: Final = urlsplit(database_url)
|
||||
host: Final = parsed.hostname
|
||||
port: Final = parsed.port
|
||||
username: Final = parsed.username
|
||||
assert host is not None and port is not None and username is not None
|
||||
extra_env: Final = {
|
||||
"DATABASE_HOST": f"{host}:{port}",
|
||||
"DATABASE_USERNAME": username,
|
||||
"DATABASE_PASSWORD": PASSWORD,
|
||||
"DATABASE_NAME": parsed.path.lstrip("/"),
|
||||
}
|
||||
completed: Final = run_v1_migrations(None, tmp_path, extra_env, request.node.name)
|
||||
output: Final = completed.output
|
||||
errors: Final = error_lines(output)
|
||||
p1000_errors: Final = tuple(line for line in errors if "P1000" in line)
|
||||
assert completed.returncode == 1, _safe_output(output)
|
||||
assert len(errors) == 4, _safe_output(output)
|
||||
assert len(p1000_errors) == 4, _safe_output(output)
|
||||
assert _retry_count_texts(errors) == (), _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_json_logs_emit_four_valid_json_p1001_error_records_without_password_fragments(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
unreachable_url: Final = _unreachable_database_url(database_url)
|
||||
completed: Final = run_v1_migrations(unreachable_url, tmp_path, {"JSON_LOGS": "true"}, request.node.name)
|
||||
output: Final = completed.output
|
||||
errors: Final = _json_error_records(output)
|
||||
messages: Final = tuple(record.get("message") for record in errors)
|
||||
assert completed.returncode == 1, _safe_output(output)
|
||||
assert len(errors) == 4, _safe_output(output)
|
||||
assert tuple(isinstance(message, str) and "P1001" in message for message in messages) == (
|
||||
True,
|
||||
True,
|
||||
True,
|
||||
True,
|
||||
), _safe_output(output)
|
||||
assert error_lines(output) == (), _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_migration_job_entrypoint_emits_four_p1001_errors_without_password_fragments(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
unreachable_url: Final = _unreachable_database_url(database_url)
|
||||
command: Final = (sys.executable, "-I", "-m", "litellm.proxy.prisma_migration")
|
||||
environment: Final = _migration_environment(
|
||||
unreachable_url,
|
||||
{"USE_V2_MIGRATION_RESOLVER": "false"},
|
||||
)
|
||||
output_path: Final = _migration_log_path(request.node.name, tmp_path)
|
||||
with output_path.open("w") as output_file:
|
||||
completed: Final = subprocess.run(
|
||||
command,
|
||||
cwd=REPO_ROOT,
|
||||
env=environment,
|
||||
stdout=output_file,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
timeout=240,
|
||||
)
|
||||
output: Final = output_path.read_text()
|
||||
errors: Final = error_lines(output)
|
||||
p1001_errors: Final = tuple(line for line in errors if "P1001" in line)
|
||||
assert completed.returncode == 1, _safe_output(output)
|
||||
assert len(errors) == 4, _safe_output(output)
|
||||
assert len(p1001_errors) == 4, _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_v2_resolver_unreachable_database_exits_2_and_names_p1001(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
unreachable_url: Final = _unreachable_database_url(database_url)
|
||||
completed: Final = run_v1_migrations(unreachable_url, tmp_path, {}, request.node.name, resolver="v2")
|
||||
output: Final = completed.output
|
||||
assert completed.returncode == 2, _safe_output(output)
|
||||
assert "P1001" in output, _safe_output(output)
|
||||
assert error_lines(output) == (), _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
||||
|
||||
def test_v2_resolver_clean_database_applies_exactly_the_shipped_migrations(
|
||||
tmp_path: Path, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
with owned_database(PASSWORD) as database_url:
|
||||
completed: Final = run_v1_migrations(database_url, tmp_path, {}, request.node.name, resolver="v2")
|
||||
output: Final = completed.output
|
||||
assert completed.returncode == 0, _safe_output(output)
|
||||
assert error_lines(output) == (), _safe_output(output)
|
||||
assert _applied_migrations(database_url) == SHIPPED_MIGRATIONS, _safe_output(output)
|
||||
_assert_no_password_fragments(output)
|
||||
|
|
@ -1,11 +1,13 @@
|
|||
import glob
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from typing import Final, NoReturn, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -19,6 +21,8 @@ sys.path.insert(
|
|||
from litellm_proxy_extras.utils import (
|
||||
PARTITIONED_SPEND_LOGS_PUSH_ERROR,
|
||||
ProxyExtrasDBManager,
|
||||
_redact_command_error,
|
||||
_redact_credentials,
|
||||
filter_partitioned_spend_logs_diff,
|
||||
)
|
||||
|
||||
|
|
@ -1412,3 +1416,411 @@ class TestMigrationJobOwnedDrift:
|
|||
assert 'PRIMARY KEY ("request_id")' not in filtered
|
||||
assert "LiteLLM_SpendLogs_legacy" not in filtered
|
||||
assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;' in filtered
|
||||
|
||||
|
||||
_P3018_UNCLASSIFIED_STDERR: Final = (
|
||||
"Error: P3018\n\n"
|
||||
"A migration failed to apply. New migrations cannot be applied before the error is "
|
||||
"recovered from.\n\n"
|
||||
"Migration name: 20260921190000_agent_identity\n\n"
|
||||
"Database error code: 23505\n\n"
|
||||
"Database error:\n"
|
||||
'ERROR: could not create unique index "agent_identity_key"\n'
|
||||
"DETAIL: Key (agent_id)=(agent-1) is duplicated.\n"
|
||||
)
|
||||
|
||||
|
||||
_FAKE_PRISMA_PID: Final = 424242
|
||||
|
||||
|
||||
class TestV1MigrationFailuresLogAtError:
|
||||
@staticmethod
|
||||
def _run_v1_migrations(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
*,
|
||||
deploy_stderr: Optional[str] = None,
|
||||
deploy_timeout: bool = False,
|
||||
diff_stderr: Optional[str] = None,
|
||||
database_url: Optional[str] = None,
|
||||
) -> tuple[bool, list[list[str]], list[int]]:
|
||||
import litellm_proxy_extras.utils as utils_module
|
||||
|
||||
calls: Final[list[list[str]]] = []
|
||||
killed_pids: Final[list[int]] = []
|
||||
|
||||
class _FakePrismaPopen:
|
||||
def __init__(
|
||||
self,
|
||||
argv: tuple[str, ...],
|
||||
*,
|
||||
env: Optional[dict[str, str]] = None,
|
||||
stdout: object = None,
|
||||
stderr: object = None,
|
||||
text: object = None,
|
||||
start_new_session: object = None,
|
||||
) -> None:
|
||||
self.args: Final = argv
|
||||
self.argv: Final = argv
|
||||
self.pid: Final = _FAKE_PRISMA_PID
|
||||
self.returncode: Optional[int] = None
|
||||
calls.append(list(argv))
|
||||
|
||||
def __enter__(self) -> "_FakePrismaPopen":
|
||||
return self
|
||||
|
||||
def __exit__(self, *args: object) -> None:
|
||||
return None
|
||||
|
||||
def _subcommand(self) -> tuple[str, str]:
|
||||
known: Final = (
|
||||
("migrate", "deploy"),
|
||||
("migrate", "diff"),
|
||||
("migrate", "resolve"),
|
||||
("db", "execute"),
|
||||
)
|
||||
for index in range(len(self.argv) - 1):
|
||||
pair: Final = tuple(self.argv[index : index + 2])
|
||||
if pair in known:
|
||||
return pair
|
||||
return ("", "")
|
||||
|
||||
def communicate(self, timeout: Optional[float] = None) -> tuple[str, str]:
|
||||
subcommand: Final = self._subcommand()
|
||||
if subcommand == ("migrate", "deploy"):
|
||||
if deploy_timeout:
|
||||
raise subprocess.TimeoutExpired(self.argv, timeout)
|
||||
if deploy_stderr is not None:
|
||||
self.returncode = 1
|
||||
return "", deploy_stderr
|
||||
self.returncode = 0
|
||||
return "No pending migrations to apply", ""
|
||||
if subcommand == ("migrate", "diff") and diff_stderr is not None:
|
||||
self.returncode = 1
|
||||
return "", diff_stderr
|
||||
self.returncode = 0
|
||||
return "", ""
|
||||
|
||||
migration_dir: Final = tmp_path / "migration_dir"
|
||||
migration_dir.mkdir()
|
||||
if database_url is None:
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("DATABASE_URL", database_url)
|
||||
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(migration_dir))
|
||||
monkeypatch.setattr(
|
||||
utils_module.prisma_toolchain.subprocess, "Popen", _FakePrismaPopen
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
utils_module.prisma_toolchain.os, "killpg", lambda pid, sig: killed_pids.append(pid)
|
||||
)
|
||||
monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
|
||||
|
||||
succeeded: Final = ProxyExtrasDBManager._run_migrations(use_migrate=True, use_v2_resolver=False)
|
||||
return succeeded, calls, killed_pids
|
||||
|
||||
@staticmethod
|
||||
def _deploy_call_count(calls: list[list[str]]) -> int:
|
||||
return sum(1 for call in calls if tuple(call[-2:]) == ("migrate", "deploy"))
|
||||
|
||||
@staticmethod
|
||||
def _error_messages(caplog: pytest.LogCaptureFixture) -> list[str]:
|
||||
return [
|
||||
record.getMessage()
|
||||
for record in caplog.records
|
||||
if record.levelno >= logging.ERROR and record.name.startswith("litellm_proxy_extras")
|
||||
]
|
||||
|
||||
def test_an_unrecognized_prisma_error_logs_its_stderr_at_error(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
stderr: Final = "Error: P1001: Can't reach database server at db:5432"
|
||||
with caplog.at_level(logging.ERROR, logger="litellm_proxy_extras"):
|
||||
succeeded, calls, _ = self._run_v1_migrations(
|
||||
monkeypatch, tmp_path, deploy_stderr=stderr
|
||||
)
|
||||
|
||||
assert succeeded is False
|
||||
assert self._deploy_call_count(calls) == 4
|
||||
assert any(stderr in message for message in self._error_messages(caplog))
|
||||
|
||||
def test_an_unclassified_p3018_logs_its_stderr_and_retry_failure_at_error(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
with caplog.at_level(logging.ERROR, logger="litellm_proxy_extras"):
|
||||
succeeded, calls, _ = self._run_v1_migrations(
|
||||
monkeypatch, tmp_path, deploy_stderr=_P3018_UNCLASSIFIED_STDERR
|
||||
)
|
||||
|
||||
assert succeeded is False
|
||||
assert self._deploy_call_count(calls) == 4
|
||||
messages: Final = self._error_messages(caplog)
|
||||
assert any(
|
||||
"20260921190000_agent_identity" in message and "is duplicated" in message
|
||||
for message in messages
|
||||
)
|
||||
assert any(
|
||||
"The process failed to execute" in message and "Retrying... (3 attempts left)" in message
|
||||
for message in messages
|
||||
)
|
||||
|
||||
def test_called_process_error_with_no_command_retries_all_v1_attempts(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
import litellm_proxy_extras.utils as utils_module
|
||||
|
||||
migration_dir: Final = tmp_path / "migration_dir"
|
||||
migration_dir.mkdir()
|
||||
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(migration_dir))
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
monkeypatch.setenv("PRISMA_OFFLINE_MODE", "true")
|
||||
monkeypatch.setenv("PRISMA_CLI_PATH", sys.executable)
|
||||
monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
|
||||
calls: Final[list[None]] = []
|
||||
|
||||
class _FakePrismaPopen:
|
||||
def __init__(
|
||||
self,
|
||||
argv: tuple[str, ...],
|
||||
*,
|
||||
env: Optional[dict[str, str]] = None,
|
||||
stdout: object = None,
|
||||
stderr: object = None,
|
||||
text: object = None,
|
||||
start_new_session: object = None,
|
||||
) -> None:
|
||||
self.args: Final = None
|
||||
self.returncode: Final = 1
|
||||
calls.append(None)
|
||||
|
||||
def __enter__(self) -> "_FakePrismaPopen":
|
||||
return self
|
||||
|
||||
def __exit__(self, *args: object) -> None:
|
||||
return None
|
||||
|
||||
def communicate(self, timeout: Optional[float] = None) -> tuple[str, str]:
|
||||
return "", "Error: P3018 unclassified"
|
||||
|
||||
monkeypatch.setattr(
|
||||
utils_module.prisma_toolchain.subprocess, "Popen", _FakePrismaPopen
|
||||
)
|
||||
|
||||
try:
|
||||
succeeded: Final = ProxyExtrasDBManager._run_migrations(
|
||||
use_migrate=True, use_v2_resolver=False
|
||||
)
|
||||
except TypeError as error:
|
||||
pytest.fail(
|
||||
f"_run_migrations raised TypeError after {len(calls)} Popen calls: {error}",
|
||||
pytrace=False,
|
||||
)
|
||||
|
||||
assert succeeded is False
|
||||
assert len(calls) == 4
|
||||
|
||||
def test_a_timeout_logs_at_error_naming_the_migrate_deploy_timeout_env_var(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
from litellm_proxy_extras.prisma_toolchain import PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="litellm_proxy_extras"):
|
||||
succeeded, calls, killed_pids = self._run_v1_migrations(
|
||||
monkeypatch, tmp_path, deploy_timeout=True
|
||||
)
|
||||
|
||||
assert succeeded is False
|
||||
assert self._deploy_call_count(calls) == 4
|
||||
assert killed_pids == [_FAKE_PRISMA_PID] * 4
|
||||
assert any(
|
||||
"timed out" in message and PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR in message
|
||||
for message in self._error_messages(caplog)
|
||||
)
|
||||
|
||||
def test_a_recovered_baseline_logs_nothing_at_error(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
with caplog.at_level(logging.ERROR, logger="litellm_proxy_extras"):
|
||||
succeeded, calls, _ = self._run_v1_migrations(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
deploy_stderr=_P3005_STDERR,
|
||||
database_url="postgresql://user:pass@db:5432/litellm",
|
||||
)
|
||||
|
||||
assert succeeded is True
|
||||
assert tuple(calls[0][-2:]) == ("migrate", "deploy")
|
||||
assert self._error_messages(caplog) == []
|
||||
|
||||
def test_a_failed_baseline_recovery_logs_its_stderr_at_error(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
database_url: Final = "postgresql://llmproxy:s3cr3t 'p\"w@db:5432/litellm"
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
with caplog.at_level(logging.DEBUG, logger="litellm_proxy_extras"):
|
||||
succeeded, calls, _ = self._run_v1_migrations(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
deploy_stderr=_P3005_STDERR,
|
||||
diff_stderr=f"baseline diff failed: XYZ-7731 for {database_url}",
|
||||
database_url=database_url,
|
||||
)
|
||||
|
||||
assert succeeded is False
|
||||
assert self._deploy_call_count(calls) == 4
|
||||
assert [
|
||||
record.getMessage()
|
||||
for record in caplog.records
|
||||
if "s3cr3t" in record.getMessage() or 'p"w' in record.getMessage()
|
||||
] == []
|
||||
messages: Final = self._error_messages(caplog)
|
||||
assert any("postgresql://REDACTED@db:5432/litellm" in message for message in messages)
|
||||
assert any("XYZ-7731" in message for message in messages)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"database_url,direct_url,text,expected",
|
||||
(
|
||||
(
|
||||
"postgresql://u:pa ss@db:5432/litellm",
|
||||
None,
|
||||
'Error: P1000: Authentication failed against database server at "postgresql://u:pa ss@db:5432/litellm"',
|
||||
'Error: P1000: Authentication failed against database server at "postgresql://REDACTED@db:5432/litellm"',
|
||||
),
|
||||
(
|
||||
"postgresql://u:pa'ss@db:5432/litellm",
|
||||
None,
|
||||
"postgresql://u:pa'ss@db:5432/litellm",
|
||||
"postgresql://REDACTED@db:5432/litellm",
|
||||
),
|
||||
(
|
||||
'postgresql://u:pa"ss@db:5432/litellm',
|
||||
None,
|
||||
'postgresql://u:pa"ss@db:5432/litellm',
|
||||
"postgresql://REDACTED@db:5432/litellm",
|
||||
),
|
||||
(
|
||||
"postgresql://u:p@ss@db:5432/litellm",
|
||||
None,
|
||||
"postgresql://u:p@ss@db:5432/litellm",
|
||||
"postgresql://REDACTED@db:5432/litellm",
|
||||
),
|
||||
(
|
||||
"postgresql://u:p%20ss@db:5432/litellm",
|
||||
None,
|
||||
"postgresql://u:p ss@db:5432/litellm",
|
||||
"postgresql://REDACTED@db:5432/litellm",
|
||||
),
|
||||
(
|
||||
"postgresql://db/litellm?password=a b&sslmode=require",
|
||||
None,
|
||||
"postgresql://db/litellm?password=a b&sslmode=require",
|
||||
"postgresql://db/litellm?REDACTED&sslmode=require",
|
||||
),
|
||||
(
|
||||
"postgresql://db/litellm?sslpassword=zq'7x",
|
||||
None,
|
||||
"postgresql://db/litellm?sslpassword=zq'7x",
|
||||
"postgresql://db/litellm?REDACTED",
|
||||
),
|
||||
(
|
||||
None,
|
||||
"postgresql://u:pa ss@db:5432/litellm",
|
||||
"postgresql://u:pa ss@db:5432/litellm",
|
||||
"postgresql://REDACTED@db:5432/litellm",
|
||||
),
|
||||
(
|
||||
None,
|
||||
None,
|
||||
"postgresql://u:pw@db/x",
|
||||
"postgresql://REDACTED@db/x",
|
||||
),
|
||||
(
|
||||
"postgresql://u:p@db:5432/litellm",
|
||||
None,
|
||||
"Error: P1001: Can't reach database server at db:5432",
|
||||
"Error: P1001: Can't reach database server at db:5432",
|
||||
),
|
||||
(
|
||||
"postgresql://u:p@db:5432/litellm",
|
||||
None,
|
||||
"Error:P1001: Can't reach database server at db:5432",
|
||||
"Error:P1001: Can't reach database server at db:5432",
|
||||
),
|
||||
(
|
||||
None,
|
||||
None,
|
||||
"plain text with no URL",
|
||||
"plain text with no URL",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_redact_credentials_masks_passwords_in_embedded_urls(
|
||||
database_url: str | None,
|
||||
direct_url: str | None,
|
||||
text: str,
|
||||
expected: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
if database_url is None:
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("DATABASE_URL", database_url)
|
||||
if direct_url is None:
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("DIRECT_URL", direct_url)
|
||||
assert _redact_credentials(text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("password", ("zq'7x", 'zq"7x', "zq'\"7x", "zq 7x", "zq@7x"))
|
||||
def test_redact_command_error_masks_url_arguments(password: str, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
database_url: Final = f"postgresql://u:{password}@db:5432/litellm"
|
||||
monkeypatch.setenv("DATABASE_URL", database_url)
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
error: Final = subprocess.CalledProcessError(1, ["prisma", "migrate", "diff", "--to-url", database_url])
|
||||
|
||||
message: Final = _redact_command_error(error)
|
||||
|
||||
assert "zq" not in message
|
||||
assert "7x" not in message
|
||||
assert "postgresql://REDACTED@db:5432/litellm" in message
|
||||
assert "returned non-zero exit status 1" in message
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"command",
|
||||
(
|
||||
None,
|
||||
Path("/usr/bin/prisma"),
|
||||
7,
|
||||
("prisma", "migrate", "deploy"),
|
||||
["prisma", "migrate", "deploy"],
|
||||
"prisma migrate deploy",
|
||||
),
|
||||
)
|
||||
def test_redact_command_error_preserves_unredacted_command_format(
|
||||
command: object, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
error: Final = subprocess.CalledProcessError(1, command)
|
||||
|
||||
assert _redact_command_error(error) == str(error)
|
||||
|
||||
|
||||
def test_redact_command_error_masks_password_in_tuple_url_argument(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
password: Final = "zq 7x"
|
||||
database_url: Final = f"postgresql://u:{password}@db:5432/litellm"
|
||||
monkeypatch.setenv("DATABASE_URL", database_url)
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
error: Final = subprocess.CalledProcessError(1, ("prisma", "migrate", "deploy", "--to-url", database_url))
|
||||
|
||||
message: Final = _redact_command_error(error)
|
||||
|
||||
assert message.startswith("Command '('")
|
||||
assert password not in message
|
||||
assert "postgresql://REDACTED@db:5432/litellm" in message
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue