Release to staging: honor DDS's upload_mode "proxy" hint in FolderIngest (ref #306834)

src/eea_datalakehouse/dds_ingestion/folder.py

@@ -32,6 +32,7 @@
     FileSpec,
     Intent,
     StatusResult,
+    UploadMode,
     UploadTarget,
 )
 from .progress import make_progress_bar
@@ -69,6 +70,7 @@ class IngestOutcome:
     files_skipped: int
     etags: dict[str, str] = field(default_factory=dict)
     resumed: bool = False
+    upload_mode: str = "direct"  # how the files went: "direct" (presigned S3) or "proxy" (DDS)
 
 
 def scan_folder(folder: Path, data_format: DataFormat) -> list[FileSpec]:
@@ -138,12 +140,17 @@ def __init__(
         idempotency_key: str | None = None,
         multipart: bool | None = None,
         show_progress: bool = True,
+        upload_mode: UploadMode = "auto",
         client: IngestClient | None = None,
         base_url: str | None = None,
         creds: DremioCreds | None = None,
     ) -> None:
         if parallelism < 1:
             raise ValueError("parallelism must be >= 1")
+        if upload_mode not in ("auto", "direct", "proxy"):
+            raise ValueError(
+                f"upload_mode must be 'auto', 'direct' or 'proxy', not {upload_mode!r}"
+            )
         self.folder = Path(folder)
         if not self.folder.is_dir():
             raise NotADirectoryError(f"not a directory: {self.folder}")
@@ -157,6 +164,7 @@ def __init__(
         self.idempotency_key = idempotency_key
         self.multipart = multipart
         self.show_progress = show_progress
+        self.upload_mode = upload_mode
 
         # Set once ``begin`` runs (or by ``attach``); the handle every session
         # command below works from.
@@ -231,8 +239,17 @@ def run(self) -> IngestOutcome:
             files_uploaded=uploaded,
             files_skipped=skipped,
             etags=etags,
+            upload_mode="proxy" if self._uses_proxy(begin) else "direct",
         )
 
+    def _uses_proxy(self, begin: BeginResult) -> bool:
+        """Whether to send the files through DDS rather than to S3 directly:
+        always with ``upload_mode="proxy"``, never with ``"direct"``, and with
+        ``"auto"`` when ``begin`` says ``upload_mode: "proxy"``."""
+        if self.upload_mode == "auto":
+            return begin.upload_mode == "proxy"
+        return self.upload_mode == "proxy"
+
     def begin(self) -> BeginResult:
         """Open the transfer and get the presigned upload targets (step 1 of 3).
 
@@ -401,6 +418,9 @@ def _upload_all(
         self, begin: BeginResult
     ) -> tuple[dict[str, str], list[dict[str, object]], int, int]:
         targets = list(begin.s3.uploads)
+        proxy = self._uses_proxy(begin)
+        if proxy:
+            logger.info("uploading through DDS (proxy), not to the presigned S3 targets")
         done = self._already_done(begin.session_id)
         pending = [t for t in targets if t.rel_path not in done]
         skipped = len(targets) - len(pending)
@@ -418,7 +438,9 @@ def _upload_all(
             if not pending:
                 return etags, multipart_etags, 0, skipped
             with ThreadPoolExecutor(max_workers=self.parallelism) as pool:
-                futures = {pool.submit(self._upload_one, t): t for t in pending}
+                futures = {
+                    pool.submit(self._upload_one, t, begin.session_id, proxy): t for t in pending
+                }
                 for future in as_completed(futures):
                     target = futures[future]
                     etag, parts = future.result()
@@ -441,15 +463,21 @@ def _upload_all(
         return etags, multipart_etags, len(pending), skipped
 
     def _upload_one(
-        self, target: UploadTarget
+        self, target: UploadTarget, session_id: str, proxy: bool = False
     ) -> tuple[str | None, list[dict[str, object]] | None]:
         """Upload one file; return ``(single_etag, multipart_parts)``.
 
-        Exactly one element is non-``None``: a single-shot upload yields the ETag
-        (or ``None`` if S3 omitted it); a multipart upload yields the per-part
-        ETag list and the single ETag is ``None``.
+        At most one element is non-``None``: a single-shot upload yields the
+        ETag (or ``None`` if S3 omitted it); a multipart upload yields the
+        per-part ETag list and the single ETag is ``None``. A proxied upload
+        (`proxy`) sends the whole file through DDS — DDS writes it to S3 with
+        its own connection, so this process never contacts the storage — and
+        yields neither: there are no S3 ETags to hand to commit.
         """
         data = (self.folder / target.rel_path).read_bytes()
+        if proxy:
+            self._client.stage(session_id, target.rel_path, data)
+            return None, None
         if target.is_multipart:
             parts = self._client.upload_file_multipart(target, data)
             return None, [dict(p) for p in parts]

src/eea_datalakehouse/dds_ingestion/models.py

@@ -12,6 +12,9 @@
 
 Intent = Literal["read_only", "editable"]
 DataFormat = Literal["parquet", "csv", "json"]
+# How FolderIngest sends the bytes: "direct" to the presigned S3 targets, or
+# "proxy" through DDS (POST /api/v1/ingest/stage); "auto" follows begin's hint.
+UploadMode = Literal["auto", "direct", "proxy"]
 
 
 @dataclass(frozen=True, slots=True)
@@ -94,20 +97,32 @@ def from_json(cls, data: dict[str, Any]) -> S3Plan:
 
 @dataclass(frozen=True, slots=True)
 class BeginResult:
-    """Response from ``POST /api/v1/ingest/begin``."""
+    """Response from ``POST /api/v1/ingest/begin``.
+
+    ``upload_mode`` is the server's hint for how to send the files:
+    ``"proxy"`` means upload through DDS (``POST /api/v1/ingest/stage``)
+    because the presigned S3 targets are not usable from the client — e.g. the
+    storage's TLS certificate doesn't match the endpoint's name. ``None`` (no
+    hint, or a server predating it) means the presigned targets, as before.
+    Read from the top level of the response, or from inside ``s3``.
+    """
 
     session_id: str
     status: str
     s3: S3Plan
     collision: Any = None
+    upload_mode: str | None = None
 
     @classmethod
     def from_json(cls, data: dict[str, Any]) -> BeginResult:
+        s3 = data["s3"]
+        mode = data.get("upload_mode") or (s3.get("upload_mode") if isinstance(s3, dict) else None)
         return cls(
             session_id=data["session_id"],
             status=data["status"],
-            s3=S3Plan.from_json(data["s3"]),
+            s3=S3Plan.from_json(s3),
             collision=data.get("collision"),
+            upload_mode=str(mode) if mode else None,
         )
 
 

tests/dds_ingestion/test_client.py

@@ -3,6 +3,7 @@
 import httpx
 import pytest
 import respx
+
 from eea_datalakehouse.dds_ingestion.client import (
     IngestApiError,
     IngestClient,
@@ -371,3 +372,26 @@ def test_a_503_from_a_proxy_is_not_called_a_storage_fault(creds: DremioCreds) ->
 
     assert not isinstance(exc.value, StorageUnavailableError)
     assert "POST /api/v1/ingest/begin" in str(exc.value)
+
+
+@respx.mock
+def test_stage_uploads_one_file_through_dds_with_bearer_auth(creds: DremioCreds) -> None:
+    route = respx.post(f"{BASE_URL}/api/v1/ingest/stage").mock(
+        return_value=httpx.Response(
+            200, json={"rel_path": "sub/a.parquet", "key": "p/sub/a.parquet", "bytes": 6}
+        )
+    )
+    with IngestClient(BASE_URL, creds) as client:
+        result = client.stage("sess-1", "sub/a.parquet", b"PAR1-a")
+
+    assert (result.rel_path, result.key, result.bytes_written) == (
+        "sub/a.parquet",
+        "p/sub/a.parquet",
+        6,
+    )
+    request = route.calls.last.request
+    assert request.headers["Authorization"] == "Bearer s3cr3t-pwd"
+    body = request.read()
+    assert b'name="session_id"' in body and b"sess-1" in body
+    assert b'name="rel_path"' in body and b"sub/a.parquet" in body
+    assert b'filename="a.parquet"' in body and b"PAR1-a" in body

tests/dds_ingestion/test_folder.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 
 import pytest
+
 from eea_datalakehouse.dds_ingestion.client import StorageUnavailableError
 from eea_datalakehouse.dds_ingestion.folder import (
     FolderIngest,
@@ -329,3 +330,103 @@ def test_run_stops_before_uploading_when_dds_cannot_write_to_s3(
     assert "administrator" in str(exc.value)
     assert client.uploaded == []
     assert client.commit_calls == []
+
+
+class ProxyHintFakeClient(FakeClient):
+    """`begin` hints ``upload_mode: "proxy"``; `stage` records server-proxied uploads."""
+
+    def __init__(self, *, hint: str | None = "proxy") -> None:
+        super().__init__()
+        self.hint = hint
+        self.staged: list[tuple[str, str, bytes]] = []
+
+    def begin(self, *, files: list[FileSpec], **kwargs: object) -> BeginResult:
+        result = super().begin(files=files, **kwargs)
+        return BeginResult(
+            session_id=result.session_id,
+            status=result.status,
+            s3=result.s3,
+            upload_mode=self.hint,
+        )
+
+    def stage(self, session_id: str, rel_path: str, data: bytes) -> object:
+        with self._lock:
+            self.staged.append((session_id, rel_path, data))
+        return object()
+
+
+def test_proxy_hint_sends_every_file_through_dds_not_s3(data_folder: Path) -> None:
+    client = ProxyHintFakeClient()
+
+    outcome = _make_ingest(data_folder, client).run()
+
+    assert client.uploaded == []  # the presigned S3 targets were never used
+    assert sorted(rel for _, rel, _ in client.staged) == ["a.parquet", "sub/b.parquet"]
+    assert {sid for sid, _, _ in client.staged} == {"sess-X"}
+    assert (data_folder / "a.parquet").read_bytes() in [d for _, _, d in client.staged]
+    assert client.commit_calls[0]["multipart_etags"] is None  # nothing from S3 to pass on
+    assert outcome.upload_mode == "proxy"
+    assert outcome.files_uploaded == 2
+
+
+def test_proxy_hint_wins_over_multipart_targets(data_folder: Path) -> None:
+    class Both(ProxyHintFakeClient, MultipartFakeClient):
+        pass
+
+    client = Both()
+
+    _make_ingest(data_folder, client, multipart=True).run()
+
+    assert client.uploaded == [] and len(client.staged) == 2
+
+
+def test_no_hint_keeps_direct_uploads(data_folder: Path) -> None:
+    client = ProxyHintFakeClient(hint=None)
+
+    outcome = _make_ingest(data_folder, client).run()
+
+    assert sorted(client.uploaded) == ["a.parquet", "sub/b.parquet"]
+    assert client.staged == []
+    assert outcome.upload_mode == "direct"
+
+
+def test_upload_mode_proxy_forces_dds_even_without_a_hint(data_folder: Path) -> None:
+    client = ProxyHintFakeClient(hint=None)
+
+    _make_ingest(data_folder, client, upload_mode="proxy").run()
+
+    assert client.uploaded == [] and len(client.staged) == 2
+
+
+def test_upload_mode_direct_ignores_the_hint(data_folder: Path) -> None:
+    client = ProxyHintFakeClient(hint="proxy")
+
+    _make_ingest(data_folder, client, upload_mode="direct").run()
+
+    assert len(client.uploaded) == 2 and client.staged == []
+
+
+def test_invalid_upload_mode_rejected(data_folder: Path) -> None:
+    with pytest.raises(ValueError, match="upload_mode must be"):
+        _make_ingest(data_folder, FakeClient(), upload_mode="s3")
+
+
+@pytest.mark.parametrize(
+    ("payload", "expected"),
+    [
+        ({"upload_mode": "proxy"}, "proxy"),
+        ({"s3_extra": {"upload_mode": "proxy"}}, "proxy"),
+        ({}, None),
+    ],
+)
+def test_begin_result_reads_the_upload_mode_hint(
+    payload: dict[str, object], expected: str | None
+) -> None:
+    s3: dict[str, object] = {"bucket": "b", "key_prefix": "p/", "uploads": []}
+    data: dict[str, object] = {"session_id": "s", "status": "open", "s3": s3}
+    if "s3_extra" in payload:
+        s3.update(payload["s3_extra"])  # type: ignore[arg-type]
+    else:
+        data.update(payload)
+
+    assert BeginResult.from_json(data).upload_mode == expected