diff --git a/litellm/integrations/mavvrik/__init__.py b/litellm/integrations/mavvrik/__init__.py index ce39976bcc0..2074004012e 100644 --- a/litellm/integrations/mavvrik/__init__.py +++ b/litellm/integrations/mavvrik/__init__.py @@ -16,8 +16,6 @@ from datetime import datetime, timedelta from datetime import timezone as _tz from typing import Optional -import polars as pl - from litellm._logging import verbose_proxy_logger from litellm.constants import MAVVRIK_MAX_FETCHED_DATA_RECORDS from litellm.integrations.mavvrik.client import Client diff --git a/litellm/integrations/mavvrik/orchestrator.py b/litellm/integrations/mavvrik/orchestrator.py index 07822b670e0..679cd12876e 100644 --- a/litellm/integrations/mavvrik/orchestrator.py +++ b/litellm/integrations/mavvrik/orchestrator.py @@ -102,8 +102,16 @@ class Orchestrator: verbose_logger.warning("Orchestrator: exporting %s → %s", start, end) for export_date in self._date_range(start, end): - await self._export(export_date) - await self._advance(export_date) + 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" + ) 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 a7c28b68d2e..d8ad4216f18 100644 --- a/litellm/integrations/mavvrik/uploader.py +++ b/litellm/integrations/mavvrik/uploader.py @@ -103,14 +103,19 @@ class Uploader: return session_uri async def _finalize_upload(self, session_uri: str, gzip_bytes: bytes) -> None: - """PUT gzip bytes to the GCS session URI to complete the bulk upload.""" + """PUT gzip bytes to the GCS session URI to complete the bulk upload. + + Sends Content-Range: bytes 0-{last}/{total} per the GCS resumable upload + spec so GCS knows the object is complete (consistent with _put_chunk final=True). + """ + total = len(gzip_bytes) + content_range = f"bytes 0-{total - 1}/{total}" if total > 0 else "bytes */0" resp = await http_request( "PUT", session_uri, headers={ "Content-Type": "application/gzip", - "Content-Encoding": "gzip", - "x-goog-resumable": "stop", + "Content-Range": content_range, }, content=gzip_bytes, timeout=120.0, diff --git a/tests/test_litellm/integrations/mavvrik/test_scheduler.py b/tests/test_litellm/integrations/mavvrik/test_scheduler.py index a8327d96b0d..03e6a121dae 100644 --- a/tests/test_litellm/integrations/mavvrik/test_scheduler.py +++ b/tests/test_litellm/integrations/mavvrik/test_scheduler.py @@ -171,8 +171,12 @@ class TestRunExportLoop: assert exported_dates == ["2026-04-09", "2026-04-10"] @pytest.mark.asyncio - async def test_skips_upload_when_no_data(self): - """When export returns 0 bytes, advance still called (date was processed).""" + async def test_does_not_advance_when_no_data(self): + """When export returns 0 bytes, marker is NOT advanced — prevents silent data loss. + + 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. + """ orc = _make_orchestrator() orc._client.register = AsyncMock(return_value="2026-04-09") @@ -188,8 +192,10 @@ class TestRunExportLoop: ): await orc.run() - # advance_marker always called — even for empty dates - orc._client.advance_marker.assert_called_once() + # advance_marker must NOT be called when 0 bytes exported + 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 async def test_reports_error_on_pipeline_failure(self):