From ff93609453284416f9848229787513ba9e8f4a4e Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 8 Feb 2024 17:37:22 -0800
Subject: [PATCH] build(schema.prisma): support direct url on prisma schema

---
 litellm/proxy/proxy_cli.py  | 25 ++++++++++++++++++-------
 litellm/proxy/schema.prisma |  1 +
 schema.prisma               |  1 +
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py
index 9a71d58f8..9f0ed8de1 100644
--- a/litellm/proxy/proxy_cli.py
+++ b/litellm/proxy/proxy_cli.py
@@ -424,14 +424,25 @@ def run_server(
             if database_url is not None and isinstance(database_url, str):
                 os.environ["DATABASE_URL"] = database_url
 
-        if os.getenv("DATABASE_URL", None) is not None:
+        if (
+            os.getenv("DATABASE_URL", None) is not None
+            or os.getenv("DIRECT_URL", None) is not None
+        ):
             try:
-                ### add connection pool + pool timeout args
-                params = {"connection_limit": 100, "pool_timeout": 60}
-                database_url = os.getenv("DATABASE_URL")
-                modified_url = append_query_params(database_url, params)
-                os.environ["DATABASE_URL"] = modified_url
-                ###
+                if os.getenv("DATABASE_URL", None) is not None:
+                    ### add connection pool + pool timeout args
+                    params = {"connection_limit": 100, "pool_timeout": 60}
+                    database_url = os.getenv("DATABASE_URL")
+                    modified_url = append_query_params(database_url, params)
+                    os.environ["DATABASE_URL"] = modified_url
+                    ###
+                if os.getenv("DIRECT_URL", None) is not None:
+                    ### add connection pool + pool timeout args
+                    params = {"connection_limit": 100, "pool_timeout": 60}
+                    database_url = os.getenv("DIRECT_URL")
+                    modified_url = append_query_params(database_url, params)
+                    os.environ["DIRECT_URL"] = modified_url
+                    ###
                 subprocess.run(["prisma"], capture_output=True)
                 is_prisma_runnable = True
             except FileNotFoundError:
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index 2164358cc..3a5c3ea23 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -1,6 +1,7 @@
 datasource client {
   provider = "postgresql"
   url      = env("DATABASE_URL")
+  directUrl = env("DIRECT_URL")
 }
 
 generator client {
diff --git a/schema.prisma b/schema.prisma
index 2164358cc..3a5c3ea23 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -1,6 +1,7 @@
 datasource client {
   provider = "postgresql"
   url      = env("DATABASE_URL")
+  directUrl = env("DIRECT_URL")
 }
 
 generator client {