build(schema.prisma): support direct url on prisma schema

BerriAI · Feb 9, 2024 · ff93609 · ff93609
1 parent f473d59
commit ff93609
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py
@@ -424,14 +424,25 @@ def _make_openai_completion():
             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
@@ -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
@@ -1,6 +1,7 @@
 datasource client {
   provider = "postgresql"
   url      = env("DATABASE_URL")
+  directUrl = env("DIRECT_URL")
 }
 
 generator client {