mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
fix(litellm_proxy_extras): add baselining db script (#9942)
* fix(litellm_proxy_extras): add baselining db script Fixes https://github.com/BerriAI/litellm/issues/9885 * fix(prisma_client.py): fix ruff errors * ci(config.yml): add publish_proxy_extras step * fix(config.yml): compare contents between versions to check for changes * fix(config.yml): fix check * fix: install toml * fix: update check * fix: ensure versions in sync * fix: fix version compare * fix: correct the cost for 'gemini/gemini-2.5-pro-preview-03-25' (#9896) * fix: Typo in the cost 'gemini/gemini-2.5-pro-preview-03-25', closes #9854 * chore: update in backup file as well * Litellm add managed files db (#9930) * fix(openai.py): ensure openai file object shows up on logs * fix(managed_files.py): return unified file id as b64 str allows retrieve file id to work as expected * fix(managed_files.py): apply decoded file id transformation * fix: add unit test for file id + decode logic * fix: initial commit for litellm_proxy support with CRUD Endpoints * fix(managed_files.py): support retrieve file operation * fix(managed_files.py): support for DELETE endpoint for files * fix(managed_files.py): retrieve file content support supports retrieve file content api from openai * fix: fix linting error * test: update tests * fix: fix linting error * feat(managed_files.py): support reading / writing files in DB * feat(managed_files.py): support deleting file from DB on delete * test: update testing * fix(spend_tracking_utils.py): ensure each file create request is logged correctly * fix(managed_files.py): fix storing / returning managed file object from cache * fix(files/main.py): pass litellm params to azure route * test: fix test * build: add new prisma migration * build: bump requirements * test: add more testing * refactor: cleanup post merge w/ main * fix: fix code qa errors * [DB / Infra] Add new column team_member_permissions (#9941) * add team_member_permissions to team table * add migration.sql file * fix poetry lock * fix prisma migrations * fix poetry lock * fix migration * ui new build * fix(factory.py): correct indentation for message index increment in ollama, This fixes bug #9822 (#9943) * fix(factory.py): correct indentation for message index increment in ollama_pt function * test: add unit tests for ollama_pt function handling various message types * ci: update test * fix: fix check * ci: see what dir looks like * ci: more checks * ci: fix filepath * ci: cleanup * ci: fix ci --------- Co-authored-by: Nilanjan De <nilanjan.de@gmail.com> Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: Dan Shaw <dan@danieljshaw.com>
This commit is contained in:
Krish Dholakia
GitHub
parent
433075a8d9
commit
d004fb542f
3 changed files with 215 additions and 89 deletions
|
|
@ -1,7 +1,9 @@
|
|||
import glob
|
||||
import os
|
||||
import random
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
|
|
@ -14,6 +16,94 @@ def str_to_bool(value: Optional[str]) -> bool:
|
|||
|
||||
|
||||
class ProxyExtrasDBManager:
|
||||
@staticmethod
|
||||
def _get_prisma_dir() -> str:
|
||||
"""Get the path to the migrations directory"""
|
||||
migrations_dir = os.path.dirname(__file__)
|
||||
return migrations_dir
|
||||
|
||||
@staticmethod
|
||||
def _create_baseline_migration(schema_path: str) -> bool:
|
||||
"""Create a baseline migration for an existing database"""
|
||||
prisma_dir = ProxyExtrasDBManager._get_prisma_dir()
|
||||
prisma_dir_path = Path(prisma_dir)
|
||||
init_dir = prisma_dir_path / "migrations" / "0_init"
|
||||
|
||||
# Create migrations/0_init directory
|
||||
init_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Generate migration SQL file
|
||||
migration_file = init_dir / "migration.sql"
|
||||
|
||||
try:
|
||||
# Generate migration diff with increased timeout
|
||||
subprocess.run(
|
||||
[
|
||||
"prisma",
|
||||
"migrate",
|
||||
"diff",
|
||||
"--from-empty",
|
||||
"--to-schema-datamodel",
|
||||
str(schema_path),
|
||||
"--script",
|
||||
],
|
||||
stdout=open(migration_file, "w"),
|
||||
check=True,
|
||||
timeout=30,
|
||||
) # 30 second timeout
|
||||
|
||||
# Mark migration as applied with increased timeout
|
||||
subprocess.run(
|
||||
[
|
||||
"prisma",
|
||||
"migrate",
|
||||
"resolve",
|
||||
"--applied",
|
||||
"0_init",
|
||||
],
|
||||
check=True,
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
return True
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(
|
||||
"Migration timed out - the database might be under heavy load."
|
||||
)
|
||||
return False
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(f"Error creating baseline migration: {e}")
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _get_migration_names(migrations_dir: str) -> list:
|
||||
"""Get all migration directory names from the migrations folder"""
|
||||
migration_paths = glob.glob(f"{migrations_dir}/migrations/*/migration.sql")
|
||||
logger.info(f"Found {len(migration_paths)} migrations at {migrations_dir}")
|
||||
return [Path(p).parent.name for p in migration_paths]
|
||||
|
||||
@staticmethod
|
||||
def _resolve_all_migrations(migrations_dir: str):
|
||||
"""Mark all existing migrations as applied"""
|
||||
migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
|
||||
logger.info(f"Resolving {len(migration_names)} migrations")
|
||||
for migration_name in migration_names:
|
||||
try:
|
||||
logger.info(f"Resolving migration: {migration_name}")
|
||||
subprocess.run(
|
||||
["prisma", "migrate", "resolve", "--applied", migration_name],
|
||||
timeout=60,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
logger.debug(f"Resolved migration: {migration_name}")
|
||||
except subprocess.CalledProcessError as e:
|
||||
if "is already recorded as applied in the database." not in e.stderr:
|
||||
logger.warning(
|
||||
f"Failed to resolve migration {migration_name}: {e.stderr}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def setup_database(schema_path: str, use_migrate: bool = False) -> bool:
|
||||
"""
|
||||
|
|
@ -30,7 +120,7 @@ class ProxyExtrasDBManager:
|
|||
use_migrate = str_to_bool(os.getenv("USE_PRISMA_MIGRATE")) or use_migrate
|
||||
for attempt in range(4):
|
||||
original_dir = os.getcwd()
|
||||
migrations_dir = os.path.dirname(__file__)
|
||||
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
|
||||
os.chdir(migrations_dir)
|
||||
|
||||
try:
|
||||
|
|
@ -55,8 +145,16 @@ class ProxyExtrasDBManager:
|
|||
"P3005" in e.stderr
|
||||
and "database schema is not empty" in e.stderr
|
||||
):
|
||||
logger.info("Error: Database schema is not empty")
|
||||
return False
|
||||
logger.info(
|
||||
"Database schema is not empty, creating baseline migration"
|
||||
)
|
||||
ProxyExtrasDBManager._create_baseline_migration(schema_path)
|
||||
logger.info(
|
||||
"Baseline migration created, resolving all migrations"
|
||||
)
|
||||
ProxyExtrasDBManager._resolve_all_migrations(migrations_dir)
|
||||
logger.info("✅ All migrations resolved.")
|
||||
return True
|
||||
else:
|
||||
# Use prisma db push with increased timeout
|
||||
subprocess.run(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue