fix(tika-worker): stream S3 blobs to Tika instead of buffering (#189)

CHANGES.md

@@ -4,6 +4,13 @@
 
 ### Fixed
 
+- The Tika worker now streams S3-tiered blobs straight from S3 into the Tika
+  request body in chunks, instead of buffering the whole blob in memory (peak
+  ~2–3× blob size). A single large PDF/image could OOM-kill a small worker pod
+  (256Mi) into `CrashLoopBackOff`; streaming keeps the worker's memory roughly
+  constant regardless of blob size. The PG-bytea path (small inline blobs below
+  the S3 tiering threshold) is unchanged. #189
+
 - The Advanced-tab *Update Catalog* and *Clear and Rebuild* buttons now submit
   via `POST` instead of `GET`. The forms in `catalogAdvanced.dtml` had no
   `method`, so the (destructive) action stayed in the URL bar and a reload /

src/plone/pgcatalog/tika_worker.py

@@ -65,6 +65,9 @@
     "    pip install plone-pgcatalog[tika-s3]   # + S3-tiered blobs"
 )
 
+# Chunk size for streaming S3 blobs into the Tika request body (#189).
+_S3_STREAM_CHUNK = 1024 * 1024  # 1 MiB
+
 
 class TikaWorker:
     """PostgreSQL-backed text extraction worker using Apache Tika."""
@@ -190,24 +193,45 @@ def _process_one(self):
             return True
 
     def _extract(self, conn, zoid, tid, content_type):
-        """Fetch blob and send to Tika, return extracted text."""
-        blob_data = self._fetch_blob(conn, zoid, tid)
+        """Fetch blob and send to Tika, return extracted text.
+
+        S3-tiered blobs are streamed straight from S3 into the Tika request
+        body in chunks, so the worker never holds the whole blob in memory
+        (#189).  PG-bytea blobs are sub-threshold (small) and sent as bytes.
+        """
         if httpx is None:
             raise RuntimeError(MISSING_EXTRA_HINT)
         headers = {"Accept": "text/plain"}
         if content_type:
             headers["Content-Type"] = content_type
+
+        source = self._blob_source(conn, zoid, tid)
         with httpx.Client(timeout=self.http_timeout) as client:
+            if source["kind"] == "s3":
+                # Stream S3 -> Tika without buffering the whole blob.  An
+                # explicit Content-Length (from get_object) lets us avoid
+                # chunked transfer-encoding.
+                obj = self._get_s3_client().get_object(
+                    Bucket=self.s3_config["bucket_name"], Key=source["s3_key"]
+                )
+                headers["Content-Length"] = str(obj["ContentLength"])
+                content = obj["Body"].iter_chunks(chunk_size=_S3_STREAM_CHUNK)
+            else:
+                content = source["data"]
             resp = client.put(
                 f"{self.tika_url}/tika",
-                content=blob_data,
+                content=content,
                 headers=headers,
             )
             resp.raise_for_status()
             return resp.text
 
-    def _fetch_blob(self, conn, zoid, tid):
-        """Fetch blob bytes from PG bytea or S3."""
+    def _blob_source(self, conn, zoid, tid):
+        """Locate a blob without materializing S3 data.
+
+        Returns ``{"kind": "bytes", "data": ...}`` for inline PG-bytea blobs
+        or ``{"kind": "s3", "s3_key": ...}`` for S3-tiered blobs.
+        """
         with conn.cursor(row_factory=dict_row) as cur:
             cur.execute(
                 "SELECT data, s3_key FROM blob_state "
@@ -218,13 +242,24 @@ def _fetch_blob(self, conn, zoid, tid):
         if row is None:
             raise ValueError(f"No blob for zoid={zoid} tid={tid}")
         if row["data"] is not None:
-            return bytes(row["data"])
+            return {"kind": "bytes", "data": bytes(row["data"])}
         if row["s3_key"]:
-            return self._fetch_from_s3(row["s3_key"])
+            return {"kind": "s3", "s3_key": row["s3_key"]}
         raise ValueError(f"Blob row has neither data nor s3_key for zoid={zoid}")
 
+    def _fetch_blob(self, conn, zoid, tid):
+        """Fetch blob bytes from PG bytea or S3 (materializes the blob).
+
+        Used by callers that need the bytes in hand; ``_extract`` streams S3
+        instead (see #189).
+        """
+        source = self._blob_source(conn, zoid, tid)
+        if source["kind"] == "bytes":
+            return source["data"]
+        return self._fetch_from_s3(source["s3_key"])
+
     def _fetch_from_s3(self, s3_key):
-        """Download blob from S3."""
+        """Download a full blob from S3 into memory."""
         import io
 
         client = self._get_s3_client()

tests/test_tika_worker_extras.py

@@ -7,6 +7,8 @@
 it must give a clear hint pointing at the extra. See issue #171.
 """
 
+from unittest import mock
+
 import subprocess
 import sys
 import textwrap
@@ -185,7 +187,9 @@ def put(self, *a, **k):
 
     monkeypatch.setattr(tika_worker, "httpx", types.SimpleNamespace(Client=_Client))
     worker = tika_worker.TikaWorker(dsn="x", tika_url="http://t", http_timeout=7.5)
-    monkeypatch.setattr(worker, "_fetch_blob", lambda *a, **k: b"blob")
+    monkeypatch.setattr(
+        worker, "_blob_source", lambda *a, **k: {"kind": "bytes", "data": b"blob"}
+    )
     out = worker._extract(conn=None, zoid=1, tid=1, content_type="application/pdf")
     assert out == "extracted text"
     assert captured["timeout"] == 7.5
@@ -219,3 +223,90 @@ def run(self):
     tika_worker.main()
 
     assert captured["http_timeout"] == 300.0
+
+
+# ---------------------------------------------------------------------------
+# #189: stream S3 -> Tika instead of buffering the whole blob in memory.
+# ---------------------------------------------------------------------------
+
+
+def _capture_httpx(monkeypatch):
+    """Patch tika_worker.httpx with a client that records the put() call."""
+    from plone.pgcatalog import tika_worker
+
+    import types
+
+    captured = {}
+
+    class _Resp:
+        text = "ok"
+
+        def raise_for_status(self):
+            pass
+
+    class _Client:
+        def __init__(self, *a, **k):
+            pass
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *a):
+            return False
+
+        def put(self, url, content=None, headers=None, **k):
+            captured["url"] = url
+            captured["content"] = content
+            captured["headers"] = headers
+            return _Resp()
+
+    monkeypatch.setattr(tika_worker, "httpx", types.SimpleNamespace(Client=_Client))
+    return captured
+
+
+def test_s3_blob_is_streamed_not_buffered(monkeypatch):
+    from plone.pgcatalog import tika_worker
+
+    captured = _capture_httpx(monkeypatch)
+    worker = tika_worker.TikaWorker(
+        dsn="x", tika_url="http://t", s3_config={"bucket_name": "b"}
+    )
+    monkeypatch.setattr(
+        worker, "_blob_source", lambda *a: {"kind": "s3", "s3_key": "k"}
+    )
+
+    chunks = iter([b"chunk1", b"chunk2"])
+
+    class _Body:
+        def iter_chunks(self, chunk_size):
+            _Body.chunk_size = chunk_size
+            return chunks
+
+    fake_s3 = mock.Mock()
+    fake_s3.get_object.return_value = {"Body": _Body(), "ContentLength": 12}
+    monkeypatch.setattr(worker, "_get_s3_client", lambda: fake_s3)
+
+    out = worker._extract(conn=None, zoid=1, tid=1, content_type="application/pdf")
+    assert out == "ok"
+    # streamed: the iterator itself is handed to httpx, not a bytes buffer
+    assert captured["content"] is chunks
+    assert not isinstance(captured["content"], (bytes, bytearray))
+    # explicit Content-Length avoids chunked transfer-encoding
+    assert captured["headers"]["Content-Length"] == "12"
+    # never buffered the whole blob
+    fake_s3.download_fileobj.assert_not_called()
+    fake_s3.get_object.assert_called_once_with(Bucket="b", Key="k")
+
+
+def test_pg_bytea_blob_sent_as_bytes(monkeypatch):
+    from plone.pgcatalog import tika_worker
+
+    captured = _capture_httpx(monkeypatch)
+    worker = tika_worker.TikaWorker(dsn="x", tika_url="http://t")
+    monkeypatch.setattr(
+        worker, "_blob_source", lambda *a: {"kind": "bytes", "data": b"PDFBYTES"}
+    )
+    out = worker._extract(conn=None, zoid=1, tid=1, content_type="application/pdf")
+    assert out == "ok"
+    assert captured["content"] == b"PDFBYTES"
+    assert "Content-Length" not in (captured["headers"] or {})