test: throwaway port of key spend update to native sqlx for fuzz check

litellm-rust/Cargo.lock

@@ -4344,6 +4344,7 @@ dependencies = [
  "litellm-cache-response",
  "litellm-callbacks-legacy-python",
  "litellm-core-utils",
+ "litellm-db",
  "litellm-error",
  "litellm-host",
  "litellm-host-python",

litellm-rust/crates/python-bridge/Cargo.toml

@@ -19,6 +19,8 @@ huggingface = ["litellm-token-counter/huggingface"]
 tiktoken = ["litellm-token-counter/tiktoken"]
 
 [dependencies]
+litellm-db.workspace = true
+sqlx.workspace = true
 litellm-secret-value.workspace = true
 fancy-regex.workspace = true
 litellm-tracing.workspace = true

litellm-rust/crates/python-bridge/src/db.rs

@@ -1,19 +1,56 @@
-use litellm_host_python::ready_future;
-use pyo3::{prelude::*, types::PyNone};
+use std::str::FromStr;
+
+use litellm_db::{EntitySpend, add_key_spend};
+use litellm_host_python::run_async_value;
+use pyo3::{exceptions::PyRuntimeError, prelude::*};
+use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
 
 #[pyclass(frozen)]
-pub(crate) struct NativeDatabase {}
+pub(crate) struct NativeDatabase {
+    pool: PgPool,
+}
+
+fn connect_options() -> PyResult<PgConnectOptions> {
+    let url = std::env::var("DATABASE_URL")
+        .map_err(|_| PyRuntimeError::new_err("DATABASE_URL is not set"))?;
+    let base = url.split('?').next().unwrap_or(&url);
+    PgConnectOptions::from_str(base)
+        .map(|options| options.application_name("litellm-native"))
+        .map_err(|error| PyRuntimeError::new_err(error.to_string()))
+}
 
-// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first
-// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens
 #[pymethods]
 impl NativeDatabase {
     #[staticmethod]
     fn connect(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
-        ready_future(py, Bound::new(py, NativeDatabase {})?.as_any())
+        let options = connect_options()?;
+        run_async_value(py, async move {
+            let pool = PgPoolOptions::new()
+                .max_connections(5)
+                .connect_lazy_with(options);
+            Ok(NativeDatabase { pool })
+        })
+    }
+
+    fn add_key_spend<'py>(
+        &self,
+        py: Python<'py>,
+        ids: Vec<String>,
+        costs: Vec<f64>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let pool = self.pool.clone();
+        run_async_value(py, async move {
+            add_key_spend(&pool, &EntitySpend { ids, costs })
+                .await
+                .map_err(|error| PyRuntimeError::new_err(error.to_string()))
+        })
     }
 
     fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        ready_future(py, PyNone::get(py).as_any())
+        let pool = self.pool.clone();
+        run_async_value(py, async move {
+            pool.close().await;
+            Ok(())
+        })
     }
 }

litellm/proxy/db/db_spend_update_writer.py

@@ -61,6 +61,7 @@ from litellm.proxy.db.daily_spend_bulk_upsert import (
     merge_by_conflict_key,
 )
 from litellm.proxy.db.db_span import db_span
+from litellm.repositories import database as database_runtime
 from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
     DailySpendUpdateQueue,
 )
@@ -2196,19 +2197,11 @@ class DBSpendUpdateWriter:
             for i in range(n_retry_times + 1):
                 start_time = time.time()
                 try:
-                    async with _spend_update_tx(prisma_client, "LiteLLM_VerificationToken") as transaction:
-                        async with transaction.batch_() as batcher:
-                            # Sort by token for consistent lock ordering across pods to prevent deadlocks.
-                            for token, response_cost in sorted(key_list_transactions.items()):
-                                spend_increment: _SpendIncrement = {"increment": response_cost}
-                                batcher.litellm_verificationtoken.update_many(  # 'update_many' prevents error from being raised if no row exists
-                                    where={"token": token},
-                                    data={
-                                        "spend": spend_increment,
-                                        "total_spend": spend_increment,
-                                        "last_active": datetime.now(timezone.utc),
-                                    },
-                                )
+                    native: Final = database_runtime.current_native
+                    if native is None:
+                        raise RuntimeError("native database is not open")
+                    ordered: Final = sorted(key_list_transactions.items())
+                    await native.add_key_spend([t for t, _ in ordered], [c for _, c in ordered])
                     break
                 except Exception as e:
                     await self.handle_spend_update_failure(

litellm/repositories/database.py

@@ -10,6 +10,9 @@ if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient
 
 
+current_native: NativeDatabase | None = None
+
+
 @dataclass(frozen=True, slots=True)
 class Database:
     prisma: "PrismaClient"
@@ -22,6 +25,8 @@ class Database:
         except BaseException:
             await prisma.aclose()
             raise
+        global current_native  # noqa: PLW0603  # fuzz port experiment: spend writer has no handle yet
+        current_native = native  # rebind-ok: fuzz port experiment
         return cls(prisma=prisma, native=native)
 
     async def aclose(self) -> None:

litellm/rust_bridge/_native.pyi

@@ -30,6 +30,7 @@ class NativeClickHouseSpendStorage:
 class NativeDatabase:
     @staticmethod
     def connect() -> Future[NativeDatabase]: ...
+    def add_key_spend(self, ids: Sequence[str], costs: Sequence[float]) -> Future[None]: ...
     def close(self) -> Future[None]: ...
 
 @final