From b320f949fa505b5e57ad15851b6e5eb324922da2 Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Fri, 24 Apr 2026 17:08:23 +0530 Subject: [PATCH] fix(mavvrik): address P0/P1/P2 issues from Greptile + CI failures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P1 — orchestrator: gate _advance on total_bytes > 0 When DB is unavailable, _export returns 0 bytes. Previously _advance was called unconditionally, permanently skipping those dates (marker advances past them). Now raises RuntimeError on 0 bytes so the single try/except catches it, calls report_error, and leaves the marker unchanged — the date is retried on the next scheduled run. P1 — __init__.py: remove top-level `import polars as pl` polars is an optional [proxy] dependency. `custom_logger_registry.py` imports Logger from this package at module level, which triggered the polars import for all SDK users. polars is not used directly in __init__.py (Exporter uses it internally) so the import can be removed entirely. P2 — uploader: add Content-Range header to _finalize_upload GCS resumable upload spec requires Content-Range: bytes 0-{last}/{total} on the final PUT. _finalize_upload was omitting it. Now consistent with _put_chunk(final=True) which already sends Content-Range correctly. CI — documentation_test_env_keys: env vars already in config_settings.md (lines 607-611), failure was from an older commit. CI — lint: factory.py Black format — pre-existing upstream issue, not our code. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- litellm/integrations/mavvrik/__init__.py | 2 -- litellm/integrations/mavvrik/orchestrator.py | 12 ++++++++++-- litellm/integrations/mavvrik/uploader.py | 11 ++++++++--- .../integrations/mavvrik/test_scheduler.py | 14 ++++++++++---- 4 files changed, 28 insertions(+), 11 deletions(-) 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):