mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
dd63637322
* test(integration): run the Langfuse DB-callback test on its own scratch database The test from #43282 wrote success_callback=langfuse and the LANGFUSE_* env into the shared integration LiteLLM_Config. The suite's long-running gateway reloads that table and only ever adds callbacks, so it kept exporting to the test's closed Langfuse fake for the rest of the shard even after the rows were restored. The owned proxy now gets a scratch database, which also removes the snapshot/restore code. scratch_database moves into _support/database.py so test_cache_and_quota and this test share one copy, and the stock-config guard now checks the callback settings instead of the raw YAML text. * test(integration): include failure_callback in the stock-config Langfuse guard
37 lines
1.4 KiB
Python
37 lines
1.4 KiB
Python
import os
|
|
import uuid
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
from typing import Final, LiteralString
|
|
from urllib.parse import urlsplit, urlunsplit
|
|
|
|
import psycopg
|
|
from psycopg import sql
|
|
from psycopg.rows import dict_row
|
|
from pydantic import JsonValue, TypeAdapter
|
|
|
|
ROWS: Final = TypeAdapter(list[dict[str, JsonValue]])
|
|
|
|
|
|
def read_rows(
|
|
query: str, parameters: tuple[str, ...], *, database_url: str | None = None
|
|
) -> list[dict[str, JsonValue]]:
|
|
with psycopg.connect(database_url or os.environ["DATABASE_URL"], row_factory=dict_row) as connection:
|
|
connection.execute("SET TRANSACTION READ ONLY")
|
|
return ROWS.validate_python(connection.execute(query, parameters).fetchall())
|
|
|
|
|
|
def write_rows(query: LiteralString, parameters: tuple[str, ...], *, database_url: str | None = None) -> None:
|
|
with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection:
|
|
connection.execute(query, parameters)
|
|
|
|
|
|
@contextmanager
|
|
def scratch_database() -> Generator[str]:
|
|
name: Final = f"integration_{uuid.uuid4().hex}"
|
|
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin:
|
|
admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name)))
|
|
try:
|
|
yield urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{name}"))
|
|
finally:
|
|
admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name)))
|