fix: Handle payload too large size and telemetry

Author: skupriienko

mailjet_rest/builders.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import base64
+import json
 import mimetypes
 import sys
 from typing import TYPE_CHECKING
@@ -182,7 +183,7 @@ def build(self) -> SendV31Message:
         if "TextPart" not in self._msg and "HTMLPart" not in self._msg and "TemplateID" not in self._msg:
             msg = "Message validation failed: TextPart, HTMLPart, or TemplateID is required."
             raise ValueError(msg)
-        if "Variables" in self._msg and sys.getsizeof(str(self._msg["Variables"])) > 1024 * 1024:
+        if "Variables" in self._msg and len(json.dumps(self._msg["Variables"])) > 1024 * 1024:
             msg = "Security Violation: Variables payload too large."
             raise ValueError(msg)
 

mailjet_rest/client.py

@@ -601,9 +601,11 @@ def _extract_telemetry(data: Any, _headers: dict[str, str] | None) -> tuple[str,
         structured_data = {}
         with suppress(Exception):
             if isinstance(data, dict):
-                # Only extract fields that are mathematically non-PII
+                messages = data.get("Messages", [{}])
+                target_dict = messages[0] if isinstance(messages, list) and messages else data
+
                 for field in _ALLOWED_TRACE_FIELDS:
-                    if val := data.get(field):
+                    if val := target_dict.get(field) or data.get(field):
                         clean_val = SecurityGuard.sanitize_log_trace(val)
                         trace_ctx.append(f"{field}={clean_val}")
                         structured_data[f"mailjet.{field.lower()}"] = clean_val

tests/unit/test_builders.py

@@ -0,0 +1,38 @@
+import pytest
+from mailjet_rest.builders import MessageBuilder
+
+
+def test_message_builder_variables_size_limit() -> None:
+    """Verify blocking of excessively large Variables objects to prevent Out-Of-Memory (OOM) errors."""
+    builder = MessageBuilder()
+
+    # Create a dictionary that exceeds 1MB when JSON-serialized
+    large_payload = {"huge_key": "x" * (1024 * 1024 + 100)}
+
+    builder._msg = {
+        "From": {"Email": "sender@example.com", "Name": "System"},
+        "To": [{"Email": "recipient@example.com"}],
+        "TextPart": "Hello",
+        "Variables": large_payload,
+    }
+
+    with pytest.raises(ValueError, match="Security Violation: Variables payload too large"):
+        builder.build()
+
+
+def test_message_builder_variables_safe_size() -> None:
+    """Verify that a valid, safe-sized Variables object passes validation without errors."""
+    builder = MessageBuilder()
+
+    builder._msg = {
+        "From": {"Email": "sender@example.com", "Name": "System"},
+        "To": [{"Email": "recipient@example.com"}],
+        "TextPart": "Hello",
+        "Variables": {"small_key": "safe_value"},
+    }
+
+    result = builder.build()
+    assert "Variables" in result
+    assert "From" in result
+    assert "To" in result
+    assert "TextPart" in result

tests/unit/test_client.py

@@ -815,3 +815,46 @@ def test_property_config_invariants(url: str, audit_flag: bool) -> None:
     # Verify configuration normalizes trailing slashes properly
     assert cfg.api_url == f"{url}/"
     assert cfg.enable_security_audit is audit_flag
+
+
+# ==========================================
+# 10. Telemetry
+# ==========================================
+
+
+def test_extract_telemetry_v3_root_level() -> None:
+    """Verify telemetry extraction from the root level of the payload (API v3)."""
+    # Use exact keys monitored within the allowed set: CustomID and TemplateID
+    payload = {"CustomID": "trace-root-123", "TemplateID": 112233}
+    trace_str, struct_data = Client._extract_telemetry(payload, None)
+
+    assert "CustomID=trace-root-123" in trace_str
+    assert "TemplateID=112233" in trace_str
+    assert struct_data["mailjet.customid"] == "trace-root-123"
+    assert struct_data["mailjet.templateid"] == "112233"
+
+def test_extract_telemetry_v31_nested_level() -> None:
+    """Verify telemetry extraction from the nested Messages array (API v3.1)."""
+    payload = {
+        "Messages": [
+            {"CustomID": "trace-nested-456", "TemplateID": 98765}
+        ]
+    }
+    trace_str, struct_data = Client._extract_telemetry(payload, None)
+
+    assert "CustomID=trace-nested-456" in trace_str
+    assert "TemplateID=98765" in trace_str
+    assert struct_data["mailjet.customid"] == "trace-nested-456"
+    assert struct_data["mailjet.templateid"] == "98765"
+
+def test_extract_telemetry_safe_fallback() -> None:
+    """Verify safe handling of invalid or empty data (no exceptions raised)."""
+    # Test empty array
+    trace_str, struct_data = Client._extract_telemetry([], None)
+    assert trace_str == ""
+    assert struct_data == {}
+
+    # Test missing Messages payload
+    trace_str, struct_data = Client._extract_telemetry({"Messages": []}, None)
+    assert trace_str == ""
+    assert struct_data == {}