fix(mavvrik): address new P1 issues from Greptile review

P1 — exporter.py: lazy polars import (SDK users)
  polars is optional [proxy] dep. Importing at module level caused SDK users
  to get ModuleNotFoundError via the import chain:
  custom_logger_registry → mavvrik/__init__ → exporter → polars.
  Now imported inside each method that uses it. TYPE_CHECKING guard for
  type annotations only.

P1 — orchestrator.py + exporter.py: distinguish DB error vs zero-traffic day
  _stream_pages now raises RuntimeError when DB is None instead of silently
  returning. This propagates through _export → caught by _run_pipeline →
  report_error, marker NOT advanced. Zero-traffic days yield nothing without
  raising → advance IS called (correct).

P1 — uploader.py: accept 200 or 201 for GCS session initiation
  Standard GCS returns 200 OK for resumable upload initiation, not 201.

P1 — test_e2e_upload.py: moved to tests/local_testing/test_mavvrik_e2e.py
  BerriAI rule: tests/test_litellm/ must contain only mock-based tests.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

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 ""

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)
 

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]}"
             )

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

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):