phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(prisma_client.py): initial commit add prisma migration support to proxy

Browse source
This commit is contained in:
Krrish Dholakia 2025-03-19 14:26:59 -07:00
parent 08c362e1b1
commit 665fdfc788
3 changed files with 156 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

33
litellm/proxy/proxy_cli.py
View file

@ -451,6 +451,12 @@ class ProxyInitializationHelpers:
help="Path to the SSL certfile. Use this when you want to provide SSL certificate when starting proxy",
envvar="SSL_CERTFILE_PATH",
)
@click.option(
"--use_prisma_migrate",
is_flag=True,
default=False,
help="Use prisma migrate instead of prisma db push for database schema updates",
)
@click.option("--local", is_flag=True, default=False, help="for local debugging")
def run_server( # noqa: PLR0915
host,
@ -486,6 +492,7 @@ def run_server( # noqa: PLR0915
ssl_keyfile_path,
ssl_certfile_path,
log_config,
use_prisma_migrate,
):
args = locals()
if local:
@ -715,7 +722,10 @@ def run_server( # noqa: PLR0915
if is_prisma_runnable:
from litellm.proxy.db.check_migration import check_prisma_schema_diff
from litellm.proxy.db.prisma_client import should_update_prisma_schema
from litellm.proxy.db.prisma_client import (
PrismaManager,
should_update_prisma_schema,
)
if (
should_update_prisma_schema(
@ -725,26 +735,7 @@ def run_server( # noqa: PLR0915
):
check_prisma_schema_diff(db_url=None)
else:
for _ in range(4):
# run prisma db push, before starting server
# Save the current working directory
original_dir = os.getcwd()
# set the working directory to where this script is
abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
os.chdir(dname)
try:
subprocess.run(
["prisma", "db", "push", "--accept-data-loss"]
)
break # Exit the loop if the subprocess succeeds
except subprocess.CalledProcessError as e:
import time
print(f"Error: {e}") # noqa
time.sleep(random.randrange(start=1, stop=5))
finally:
os.chdir(original_dir)
PrismaManager.setup_database(use_migrate=use_prisma_migrate)
else:
print( # noqa
f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa