refactor(proxy/db): type Prisma wrappers against the generated client

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

litellm/proxy/db/prisma_client.py

@@ -12,11 +12,14 @@ import urllib
 import urllib.parse
 from dataclasses import dataclass
 from datetime import datetime, timedelta
-from typing import Any, Callable, Union
+from typing import TYPE_CHECKING, Any, Callable, Union
 
 from litellm._logging import verbose_proxy_logger
 from litellm.secret_managers.main import str_to_bool
 
+if TYPE_CHECKING:
+    from prisma import Prisma
+
 
 @dataclass(frozen=True)
 class IAMEndpoint:
@@ -90,7 +93,7 @@ class PrismaWrapper:
 
     def __init__(
         self,
-        original_prisma: Any,
+        original_prisma: "Prisma",
         iam_token_db_auth: bool,
         *,
         db_url_env_var: str = "DATABASE_URL",
@@ -139,8 +142,7 @@ class PrismaWrapper:
     def _get_engine_pid(self) -> int:
         """Get the PID of the current Prisma engine subprocess, or 0 if unavailable."""
         try:
-            engine = self._original_prisma._engine
-            process = getattr(engine, "process", None) if engine is not None else None
+            process = getattr(self._original_prisma._engine, "process", None)
             if process is not None:
                 pid = process.pid
                 if isinstance(pid, int):

litellm/proxy/db/routing_prisma_wrapper.py

@@ -42,8 +42,8 @@ class _RoutedActions:
 
     def __init__(
         self,
-        writer_actions: Any,
-        reader_actions: Any,
+        writer_actions: object,
+        reader_actions: object,
         should_use_reader: Callable[[], bool],
     ):
         self._writer_actions = writer_actions
@@ -97,11 +97,11 @@ class RoutingPrismaWrapper:
     def _should_use_reader(self) -> bool:
         return not self._reader_unavailable
 
-    async def connect(self, *args: Any, **kwargs: Any) -> None:
-        await self._writer.connect(*args, **kwargs)
+    async def connect(self) -> None:
+        await self._writer.connect()
         verbose_proxy_logger.info("[writer] DB connected")
         try:
-            await self._reader.connect(*args, **kwargs)
+            await self._reader.connect()
             self._reader_unavailable = False
             verbose_proxy_logger.info("[reader] DB connected")
         except Exception as e:
@@ -116,11 +116,11 @@ class RoutingPrismaWrapper:
                 e,
             )
 
-    async def disconnect(self, *args: Any, **kwargs: Any) -> None:
+    async def disconnect(self) -> None:
         first_error: BaseException | None = None
         for client in (self._writer, self._reader):
             try:
-                await client.disconnect(*args, **kwargs)
+                await client.disconnect()
             except Exception as e:
                 if first_error is None:
                     first_error = e