diff --git a/litellm/integrations/mavvrik/exporter.py b/litellm/integrations/mavvrik/exporter.py index 8b1d3502bb1..7e49c7e9fa7 100644 --- a/litellm/integrations/mavvrik/exporter.py +++ b/litellm/integrations/mavvrik/exporter.py @@ -17,15 +17,19 @@ Internal methods: DB not connected: all methods log a warning and return empty/None — never raise. The scheduler skips the date gracefully; user-triggered endpoints surface the missing-DB error through Settings._ensure_prisma_client() before reaching here. + +polars is an optional [proxy] dependency — imported lazily inside methods so +SDK-only users are not affected when Logger is imported via custom_logger_registry. """ import io -from typing import Any, AsyncIterator, List, Optional, Tuple - -import polars as pl +from typing import Any, AsyncIterator, List, Optional, Tuple, TYPE_CHECKING from litellm._logging import verbose_proxy_logger +if TYPE_CHECKING: + import polars as pl + # query_raw is used here instead of Prisma model methods because the query # requires a 4-table LEFT JOIN (DailyUserSpend → VerificationToken → # TeamTable → UserTable). Prisma's relational API cannot express a multi-hop @@ -79,7 +83,7 @@ class Exporter: date_str: str, connection_id: Optional[str] = None, limit: Optional[int] = None, - ) -> Tuple[pl.DataFrame, str]: + ) -> Tuple["pl.DataFrame", str]: """Fetch and serialize spend data for one calendar date. All rows are exported — including failed requests. Mavvrik decides @@ -101,14 +105,19 @@ class Exporter: Uses LIMIT/OFFSET pagination so only page_size rows are in memory at once. All rows exported — including failed requests. - Yields nothing when DB is not connected or no rows exist for the date. + + When DB is not connected, raises RuntimeError so the caller (Orchestrator) + knows the export failed — distinct from a legitimate zero-traffic day which + yields nothing without raising. """ + import polars as pl + client = self._prisma_client if client is None: - verbose_proxy_logger.warning( - "Exporter: database not connected, skipping stream for %s", date_str + raise RuntimeError( + "Exporter: database not connected — cannot stream pages for " + f"{date_str}. Connect a database to your proxy." ) - return header_written = False offset = 0 @@ -122,7 +131,7 @@ class Exporter: ) if not rows: - break + break # legitimate end — no more rows (or zero-traffic day) df = pl.DataFrame(rows, infer_schema_length=None) @@ -168,11 +177,13 @@ class Exporter: self, date_str: str, limit: Optional[int] = None, - ) -> pl.DataFrame: + ) -> "pl.DataFrame": """Retrieve all spend rows for a single calendar date. Returns empty DataFrame when DB is not connected. """ + import polars as pl + client = self._prisma_client if client is None: verbose_proxy_logger.warning( @@ -191,8 +202,10 @@ class Exporter: db_response = await client.db.query_raw(query, *params) return pl.DataFrame(db_response, infer_schema_length=None) - def _to_csv(self, df: pl.DataFrame, connection_id: Optional[str] = None) -> str: + def _to_csv(self, df: "pl.DataFrame", connection_id: Optional[str] = None) -> str: """Serialize a DataFrame to CSV, adding connection_id column if provided.""" + import polars as pl + if df.is_empty(): verbose_proxy_logger.debug("Exporter: empty DataFrame, nothing to export") return "" diff --git a/litellm/integrations/mavvrik/orchestrator.py b/litellm/integrations/mavvrik/orchestrator.py index 679cd12876e..8595e971071 100644 --- a/litellm/integrations/mavvrik/orchestrator.py +++ b/litellm/integrations/mavvrik/orchestrator.py @@ -102,16 +102,11 @@ class Orchestrator: verbose_logger.warning("Orchestrator: exporting %s → %s", start, end) for export_date in self._date_range(start, end): - total_bytes = await self._export(export_date) - if total_bytes > 0: - await self._advance(export_date) - else: - # DB unavailable or no data — raise so the marker does NOT advance. - # Next run will retry this date from the same marker position. - raise RuntimeError( - f"No data streamed for {export_date.isoformat()} " - f"— DB may be unavailable; marker not advanced" - ) + # _export raises if the DB is unavailable (propagates through + # _stream_pages → _stream_upload). A genuine zero-traffic day + # returns 0 bytes without raising — advance the marker normally. + await self._export(export_date) + await self._advance(export_date) verbose_logger.warning("Orchestrator: export complete, last date=%s", end) diff --git a/litellm/integrations/mavvrik/uploader.py b/litellm/integrations/mavvrik/uploader.py index d8ad4216f18..342c875fa84 100644 --- a/litellm/integrations/mavvrik/uploader.py +++ b/litellm/integrations/mavvrik/uploader.py @@ -93,7 +93,7 @@ class Uploader: timeout=30.0, label="initiate", ) - if resp.status_code != 201: + if resp.status_code not in (200, 201): raise RuntimeError( f"GCS initiate upload failed: {resp.status_code} {resp.text[:200]}" ) diff --git a/tests/test_litellm/integrations/mavvrik/test_e2e_upload.py b/tests/local_testing/test_mavvrik_e2e.py similarity index 100% rename from tests/test_litellm/integrations/mavvrik/test_e2e_upload.py rename to tests/local_testing/test_mavvrik_e2e.py diff --git a/tests/test_litellm/integrations/mavvrik/test_scheduler.py b/tests/test_litellm/integrations/mavvrik/test_scheduler.py index 03e6a121dae..b9722210174 100644 --- a/tests/test_litellm/integrations/mavvrik/test_scheduler.py +++ b/tests/test_litellm/integrations/mavvrik/test_scheduler.py @@ -171,11 +171,13 @@ class TestRunExportLoop: assert exported_dates == ["2026-04-09", "2026-04-10"] @pytest.mark.asyncio - async def test_does_not_advance_when_no_data(self): - """When export returns 0 bytes, marker is NOT advanced — prevents silent data loss. + async def test_advances_marker_on_zero_traffic_day(self): + """Zero-traffic days (0 bytes, no DB error) still advance the marker. - 0 bytes means DB was unavailable or no data for that date. - Raising ensures the marker stays put so the date is retried next run. + A legitimate date with no spend rows returns 0 bytes without raising. + The marker must advance so the pipeline doesn't stall on quiet days. + DB-unavailability is handled differently — _stream_pages raises, which + propagates through _export and aborts before _advance is reached. """ orc = _make_orchestrator() @@ -192,9 +194,32 @@ class TestRunExportLoop: ): await orc.run() - # advance_marker must NOT be called when 0 bytes exported + # advance_marker IS called — zero bytes is a legitimate empty day + orc._client.advance_marker.assert_called_once() + orc._client.report_error.assert_not_called() + + @pytest.mark.asyncio + async def test_does_not_advance_when_db_unavailable(self): + """When DB is unavailable, _export raises — marker NOT advanced.""" + orc = _make_orchestrator() + + orc._client.register = AsyncMock(return_value="2026-04-09") + orc._client.advance_marker = AsyncMock() + orc._client.report_error = AsyncMock() + + with patch.object( + orc, + "_export", + new_callable=AsyncMock, + side_effect=RuntimeError("database not connected"), + ), patch.object( + Orchestrator, "_utc_today", return_value=date(2026, 4, 10) + ), patch.object( + Orchestrator, "_get_pod_lock_manager", return_value=None + ): + await orc.run() + orc._client.advance_marker.assert_not_called() - # error reported to Mavvrik so the failure is visible orc._client.report_error.assert_called_once() @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/mavvrik/test_transform.py b/tests/test_litellm/integrations/mavvrik/test_transform.py index 7e4687fa742..ba197bef867 100644 --- a/tests/test_litellm/integrations/mavvrik/test_transform.py +++ b/tests/test_litellm/integrations/mavvrik/test_transform.py @@ -444,18 +444,21 @@ class TestExporterNoDb: assert result is None @pytest.mark.asyncio - async def test_stream_pages_yields_nothing_when_no_db(self): - """_stream_pages yields nothing when DB not connected.""" + async def test_stream_pages_raises_when_no_db(self): + """_stream_pages raises RuntimeError when DB not connected. + + Raising (not silently returning) ensures the Orchestrator's try/except + catches it and does NOT advance the marker — preventing silent data loss. + """ exporter = Exporter() with patch.object( type(exporter), "_prisma_client", new_callable=lambda: property(lambda self: None), ): - chunks = [] - async for chunk in exporter._stream_pages("2026-04-10", connection_id="c"): - chunks.append(chunk) - assert chunks == [] + with pytest.raises(RuntimeError, match="database not connected"): + async for _ in exporter._stream_pages("2026-04-10", connection_id="c"): + pass @pytest.mark.asyncio async def test_stream_pages_yields_nothing_when_db_empty(self):