diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index a922a36d71..d95055287a 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -47,6 +47,18 @@ We structure this changelog in accordance with [Keep a Changelog](https://keepac ### Fixed +- ETL webservers now classify direct-put responses by the new + `Ais-Direct-Put-Complete` marker header, which the AIS target returns + (alongside `204` and `Ais-Direct-Put-Length`) after storing an object via + the direct-PUT endpoint. A marked ack is passed back through the pipeline + as-is: the 204, the marker, and the length (even when the length is 0 for + an empty stored object). Markerless responses use a `Content-Length`-keyed + fallback — `0` means delivered, while absent (chunked) or non-zero means + transformed content, forwarded as-is, so a chunked 200 with an empty body + (a valid empty transform result) is no longer misreported as delivered. + The fallback exists for targets that predate the marker and will be phased + out with them. `ETLServer.handle_direct_put_response` now returns a + 4-tuple `(status, body, direct_put_length, direct_put_complete)`. - ETL webservers now forward `etl_args` to the next stage on direct-put pipeline hops. Previously only the first pipeline stage received `etl_args`; stages 2..N saw an empty value. diff --git a/python/aistore/sdk/const.py b/python/aistore/sdk/const.py index 57177c4b46..3fd376b214 100644 --- a/python/aistore/sdk/const.py +++ b/python/aistore/sdk/const.py @@ -61,6 +61,11 @@ HEADER_OBJECT_BLOB_WORKERS = HEADER_PREFIX + "Blob-Workers" HEADER_OBJECT_APPEND_HANDLE = HEADER_PREFIX + "Append-Handle" HEADER_DIRECT_PUT_LENGTH = HEADER_PREFIX + "Direct-Put-Length" +# Set by the AIS target (alongside a 204 and HEADER_DIRECT_PUT_LENGTH) after a +# successful direct PUT to signal that the destination already stored the +# object. Presence-based (value ignored); intermediate ETL webservers propagate +# it back through the pipeline. +HEADER_DIRECT_PUT_COMPLETE = HEADER_PREFIX + "Direct-Put-Complete" # ETL → AIS retry contract: emitted by the ETL webserver alongside a 503 # response to signal that AIS should retry the whole PUT (the ETL bailed on # a transient direct-put failure and the request body was one-shot). diff --git a/python/aistore/sdk/etl/webserver/base_etl_server.py b/python/aistore/sdk/etl/webserver/base_etl_server.py index 5a26c1c46f..5c2ee53764 100644 --- a/python/aistore/sdk/etl/webserver/base_etl_server.py +++ b/python/aistore/sdk/etl/webserver/base_etl_server.py @@ -16,6 +16,8 @@ STATUS_OK, STATUS_BAD_GATEWAY, HEADER_AUTHORIZATION, + HEADER_CONTENT_LENGTH, + HEADER_DIRECT_PUT_COMPLETE, HEADER_DIRECT_PUT_LENGTH, AIS_AUTHN_TOKEN, AIS_DIRECT_PUT_RETRIES, @@ -64,13 +66,14 @@ def _is_connection_refused(exc: requests.ConnectionError) -> bool: def _handle_direct_put_transient_error( direct_put_url: str, exc: Exception, logger: logging.Logger -) -> Tuple[int, bytes, int]: +) -> Tuple[int, bytes, int, bool]: """ Handle a caught SYNC_DIRECT_PUT_TRANSIENT_ERRORS exception. - Returns a ``(STATUS_BAD_GATEWAY, error_bytes, 0)`` tuple for permanent - ``ConnectionRefused`` errors. Re-raises all other transient errors as - ``ETLDirectPutTransientError`` so the caller's retry loop can act on them. + Returns a ``(STATUS_BAD_GATEWAY, error_bytes, 0, False)`` tuple for + permanent ``ConnectionRefused`` errors. Re-raises all other transient + errors as ``ETLDirectPutTransientError`` so the caller's retry loop can + act on them. Args: direct_put_url: The direct-put URL that was being contacted. @@ -78,7 +81,8 @@ def _handle_direct_put_transient_error( logger: Logger used to emit the permanent-error message. Returns: - ``(STATUS_BAD_GATEWAY, encoded_error_message, 0)`` for permanent errors. + ``(STATUS_BAD_GATEWAY, encoded_error_message, 0, False)`` for + permanent errors. Raises: ETLDirectPutTransientError: For all other transient errors. @@ -86,7 +90,7 @@ def _handle_direct_put_transient_error( if isinstance(exc, requests.ConnectionError) and _is_connection_refused(exc): error = f"direct_put to {direct_put_url!r} failed: {type(exc).__name__}: {exc}".encode() logger.error("Permanent connection error to %s: %s", direct_put_url, exc) - return STATUS_BAD_GATEWAY, error, 0 + return STATUS_BAD_GATEWAY, error, 0, False raise ETLDirectPutTransientError(direct_put_url, exc) from exc @@ -259,8 +263,17 @@ def iter_and_close(output_iter: Iterator[bytes], reader) -> Iterator[bytes]: ETLServer.close_reader(reader) @staticmethod - def make_direct_put_headers(direct_put_length: int) -> dict: - """Build response headers for a direct-put result.""" + def make_direct_put_headers(direct_put_length: int, complete: bool = False) -> dict: + """Build response headers for a direct-put result. + + When `complete` (the delivered ack carried `HEADER_DIRECT_PUT_COMPLETE`), + propagate the marker and the length verbatim, including a length of 0. + """ + if complete: + return { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(direct_put_length), + } if direct_put_length != 0: return {HEADER_DIRECT_PUT_LENGTH: str(direct_put_length)} return {} @@ -293,31 +306,62 @@ def client_put( def handle_direct_put_response( self, resp: requests.Response, data: bytes, data_length: int = -1 - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """Handle the response from a direct PUT request. + Returns a `(status, body, direct_put_length, direct_put_complete)` + tuple. `direct_put_complete` is True only when the response carried + `HEADER_DIRECT_PUT_COMPLETE` — the target's ack that it stored the + object — and must be propagated (see `make_direct_put_headers`). + Args: resp: The HTTP response from the direct PUT. data: The original data bytes (used to compute length for the - 200-OK-empty-content case). Can be `b""` for streaming. + legacy 200 + `Content-Length: 0` delivered case). Can be + `b""` for streaming. data_length: Explicit byte count override. When >= 0, used instead of `len(data)`. Pass this from a `CountingIterator` for streaming pipeline PUTs where `data` is empty. """ size = data_length if data_length >= 0 else len(data) + # Delivered ack from the target (directly or propagated by a + # downstream stage). Presence-based: the value is ignored, and the + # marker decides regardless of status (the target pairs it with 204). + if HEADER_DIRECT_PUT_COMPLETE in resp.headers: + return ( + STATUS_NO_CONTENT, + b"", + int(resp.headers.get(HEADER_DIRECT_PUT_LENGTH, "0")), + True, + ) + + # Legacy handling below, unchanged: kept for targets that predate + # HEADER_DIRECT_PUT_COMPLETE; to be phased out with them. if resp.status_code == STATUS_NO_CONTENT: return ( resp.status_code, b"", int(resp.headers.get(HEADER_DIRECT_PUT_LENGTH, "0")), + False, ) if resp.status_code == STATUS_OK: - if resp.content: # from other ETL server, forward the content back - return resp.status_code, resp.content, 0 + # Keyed on the Content-Length header, mirroring the Go webserver's + # directPut (ext/etl/webserver/webserver.go): `0` means the next + # hop was the target — delivered, no content. Absent (chunked) or + # > 0 means transformed content from another ETL server; forward + # it as-is — an empty chunked body is a valid empty object. + content_length = resp.headers.get(HEADER_CONTENT_LENGTH) + try: + delivered = content_length is not None and int(content_length) == 0 + except (TypeError, ValueError): + delivered = False # malformed header: treat as content + if delivered: + return STATUS_NO_CONTENT, b"", size, False # from target, no content - return STATUS_NO_CONTENT, b"", size # from target, no content + # from other ETL server, forward the content back + return resp.status_code, resp.content, 0, False error = resp.content self.logger.error( @@ -326,7 +370,7 @@ def handle_direct_put_response( resp.status_code, error, ) - return resp.status_code, error, 0 + return resp.status_code, error, 0, False class CountingIterator: # pylint: disable=too-few-public-methods diff --git a/python/aistore/sdk/etl/webserver/fastapi_server.py b/python/aistore/sdk/etl/webserver/fastapi_server.py index f18a535797..f0cef2149b 100644 --- a/python/aistore/sdk/etl/webserver/fastapi_server.py +++ b/python/aistore/sdk/etl/webserver/fastapi_server.py @@ -50,7 +50,6 @@ ETL_WS_FQN, ETL_WS_PATH, ETL_WS_PIPELINE, - HEADER_DIRECT_PUT_LENGTH, HEADER_ETL_RETRY_REASON, ETL_RETRY_REASON_DIRECT_PUT_TRANSIENT, QPARAM_ETL_ARGS, @@ -268,7 +267,7 @@ async def _handle_request_buffered(self, path: str, request: Request, is_get: bo if pipeline_header: first_url, remaining_pipeline = parse_etl_pipeline(pipeline_header) if first_url: - status_code, transformed, direct_put_length = ( + status_code, transformed, direct_put_length, direct_put_complete = ( await self._direct_put_with_retry( first_url, transformed, remaining_pipeline, path, etl_args ) @@ -278,10 +277,8 @@ async def _handle_request_buffered(self, path: str, request: Request, is_get: bo return Response( content=transformed, status_code=status_code, - headers=( - {HEADER_DIRECT_PUT_LENGTH: str(direct_put_length)} - if direct_put_length != 0 - else {} + headers=self.make_direct_put_headers( + direct_put_length, direct_put_complete ), ) @@ -314,7 +311,7 @@ async def _handle_request_streaming( return Response( content=result[1], status_code=result[0], - headers=self.make_direct_put_headers(result[2]), + headers=self.make_direct_put_headers(result[2], result[3]), ) reader = await self._get_stream_reader(fqn, path, request, is_get) @@ -374,7 +371,7 @@ async def _direct_put_stream_with_retry( # pylint: disable=too-many-arguments,t etl_args: str, first_url: str, remaining: str, - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """ Stream-put with exponential-backoff retry on transient network errors. @@ -401,8 +398,8 @@ async def _direct_put_stream_with_retry( # pylint: disable=too-many-arguments,t forwarded to the next stage via the `AIS-Node-Url` header. Returns: - Tuple[int, bytes, int]: `(status_code, body, length)` — see - `_direct_put_stream` for semantics. + Tuple[int, bytes, int, bool]: `(status_code, body, length, + complete)` — see `_direct_put_stream` for semantics. Raises: ETLDirectPutTransientError: if all retry attempts are exhausted. @@ -458,15 +455,16 @@ async def _direct_put_stream( # pylint: disable=too-many-arguments,too-many-pos remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """ Stream transformed output directly to the next pipeline stage. Returns: - (status_code, body, length) where: + (status_code, body, length, complete) where: - status_code: HTTP status of the PUT (200/204 on success, 500 on error). - body: response bytes forwarded back to the AIS target (empty on success). - length: bytes sent to the destination, from CountingIterator. + - complete: the ack carried HEADER_DIRECT_PUT_COMPLETE; propagate it. """ try: url = compose_etl_direct_put_url( @@ -493,7 +491,7 @@ async def _direct_put_stream( # pylint: disable=too-many-arguments,too-many-pos exc, exc_info=True, ) - return STATUS_INTERNAL_SERVER_ERROR, repr(exc).encode(), 0 + return STATUS_INTERNAL_SERVER_ERROR, repr(exc).encode(), 0, False async def _get_fqn_content(self, path: str) -> bytes: """Safely read local file content with path normalization.""" @@ -531,12 +529,12 @@ async def _direct_put_with_retry( # pylint: disable=too-many-arguments,too-many remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """ Buffered direct-put with exponential-backoff retry on transient network errors. Returns: - (status_code, body, length) — see _direct_put for semantics. + (status_code, body, length, complete) — see _direct_put for semantics. Raises: ETLDirectPutTransientError: if all retry attempts are exhausted. """ @@ -565,7 +563,7 @@ async def _direct_put( # pylint: disable=too-many-arguments,too-many-positional remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """ Sends the transformed object directly to the specified AIS node (`direct_put_url`), eliminating the additional network hop through the original target. @@ -578,7 +576,8 @@ async def _direct_put( # pylint: disable=too-many-arguments,too-many-positional path: The path of the object. etl_args: Per-request transform arguments to forward to the next stage. Returns: - status code, transformed data, length of the transformed data (if any) + status code, transformed data, length of the transformed data (if any), + and whether the ack carried HEADER_DIRECT_PUT_COMPLETE Raises: ETLDirectPutTransientError: on ReadError/ConnectError/RemoteProtocolError so the caller can retry without re-fetching data. @@ -606,7 +605,7 @@ async def _direct_put( # pylint: disable=too-many-arguments,too-many-positional exc, exc_info=True, ) - return STATUS_INTERNAL_SERVER_ERROR, repr(exc).encode(), 0 + return STATUS_INTERNAL_SERVER_ERROR, repr(exc).encode(), 0, False def _build_response(self, content: bytes, mime_type: str) -> Response: """Construct standardized response with appropriate headers.""" @@ -648,7 +647,7 @@ async def _handle_ws_message(self, websocket: WebSocket): self.logger.debug("pipeline_header: %r", pipeline_header) first_url, remaining_pipeline = parse_etl_pipeline(pipeline_header) if first_url: - status_code, transformed, direct_put_length = ( + status_code, transformed, direct_put_length, _ = ( await self._direct_put_with_retry( first_url, transformed, remaining_pipeline, path, etl_args ) diff --git a/python/aistore/sdk/etl/webserver/flask_server.py b/python/aistore/sdk/etl/webserver/flask_server.py index ce6a2d50ec..4fd708e01e 100644 --- a/python/aistore/sdk/etl/webserver/flask_server.py +++ b/python/aistore/sdk/etl/webserver/flask_server.py @@ -35,7 +35,6 @@ STATUS_SERVICE_UNAVAILABLE, QPARAM_ETL_ARGS, QPARAM_ETL_FQN, - HEADER_DIRECT_PUT_LENGTH, HEADER_ETL_RETRY_REASON, ETL_RETRY_REASON_DIRECT_PUT_TRANSIENT, STATUS_INTERNAL_SERVER_ERROR, @@ -165,7 +164,7 @@ def _handle_request_buffered(self, path): if pipeline_header: first_url, remaining_pipeline = parse_etl_pipeline(pipeline_header) if first_url: - status_code, transformed, direct_put_length = ( + status_code, transformed, direct_put_length, direct_put_complete = ( self._direct_put_with_retry( first_url, transformed, remaining_pipeline, path, etl_args ) @@ -173,10 +172,8 @@ def _handle_request_buffered(self, path): return Response( response=transformed, status=status_code, - headers=( - {HEADER_DIRECT_PUT_LENGTH: str(direct_put_length)} - if direct_put_length != 0 - else {} + headers=self.make_direct_put_headers( + direct_put_length, direct_put_complete ), ) @@ -200,7 +197,7 @@ def _handle_request_streaming(self, path): return Response( response=result[1], status=result[0], - headers=self.make_direct_put_headers(result[2]), + headers=self.make_direct_put_headers(result[2], result[3]), ) # No pipeline: stream directly to client @@ -247,7 +244,7 @@ def _direct_put_with_retry( # pylint: disable=too-many-arguments,too-many-posit remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """Buffered direct-put with exponential-backoff retry on transient errors.""" for attempt in range(self.direct_put_retries + 1): try: @@ -275,7 +272,7 @@ def _direct_put_stream_with_retry( path: str, remaining_pipeline: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """ Streaming direct-put with exponential-backoff retry on transient errors. @@ -334,7 +331,7 @@ def _direct_put_stream( # pylint: disable=too-many-arguments,too-many-positiona remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """Stream transformed output directly to the next pipeline stage.""" try: url = compose_etl_direct_put_url( @@ -362,7 +359,7 @@ def _direct_put_stream( # pylint: disable=too-many-arguments,too-many-positiona root, exc_info=True, ) - return STATUS_INTERNAL_SERVER_ERROR, str(e).encode(), 0 + return STATUS_INTERNAL_SERVER_ERROR, str(e).encode(), 0, False def _handle_get(self, path): etl_args = request.args.get(QPARAM_ETL_ARGS, "").strip() @@ -404,7 +401,7 @@ def _direct_put( # pylint: disable=too-many-arguments,too-many-positional-argum remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """ Sends the transformed object directly to the specified AIS node (`direct_put_url`), eliminating the additional network hop through the original target. @@ -417,7 +414,8 @@ def _direct_put( # pylint: disable=too-many-arguments,too-many-positional-argum path: The path of the object. etl_args: Per-request transform arguments to forward to the next stage. Returns: - status code, transformed data, length of the transformed data (if any) + status code, transformed data, length of the transformed data (if any), + and whether the ack carried HEADER_DIRECT_PUT_COMPLETE """ try: url = compose_etl_direct_put_url( @@ -435,7 +433,7 @@ def _direct_put( # pylint: disable=too-many-arguments,too-many-positional-argum except Exception as e: error = str(e).encode() self.logger.error("Exception in direct put to %s: %s", direct_put_url, e) - return STATUS_INTERNAL_SERVER_ERROR, error, 0 + return STATUS_INTERNAL_SERVER_ERROR, error, 0, False # Example Gunicorn command to run this server: # command: ["gunicorn", "your_module:flask_app", "--bind", "0.0.0.0:8000", "--workers", "4"] diff --git a/python/aistore/sdk/etl/webserver/http_multi_threaded_server.py b/python/aistore/sdk/etl/webserver/http_multi_threaded_server.py index 3ac4a79eb8..44e18b2061 100644 --- a/python/aistore/sdk/etl/webserver/http_multi_threaded_server.py +++ b/python/aistore/sdk/etl/webserver/http_multi_threaded_server.py @@ -33,6 +33,7 @@ HEADER_CONTENT_LENGTH, HEADER_CONTENT_TYPE, HEADER_NODE_URL, + HEADER_DIRECT_PUT_COMPLETE, HEADER_DIRECT_PUT_LENGTH, HEADER_ETL_RETRY_REASON, ETL_RETRY_REASON_DIRECT_PUT_TRANSIENT, @@ -134,12 +135,18 @@ def _set_headers( status_code: int = STATUS_OK, length: int = 0, direct_put_length: int = 0, + direct_put_complete: bool = False, ): self.send_response(status_code) mime_type = self.server.etl_server.get_mime_type() self.send_header(HEADER_CONTENT_TYPE, mime_type) self.send_header(HEADER_CONTENT_LENGTH, str(length)) - if direct_put_length != 0: + if direct_put_complete: + # delivered ack: propagate the marker and the length verbatim, + # including a length of 0 (see make_direct_put_headers) + self.send_header(HEADER_DIRECT_PUT_COMPLETE, "true") + self.send_header(HEADER_DIRECT_PUT_LENGTH, str(direct_put_length)) + elif direct_put_length != 0: self.send_header(HEADER_DIRECT_PUT_LENGTH, str(direct_put_length)) self.end_headers() @@ -172,7 +179,7 @@ def _direct_put( # pylint: disable=too-many-arguments,too-many-positional-argum remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """ Sends the transformed object directly to the specified AIS node (`direct_put_url`), eliminating the additional network hop through the original target. @@ -185,7 +192,9 @@ def _direct_put( # pylint: disable=too-many-arguments,too-many-positional-argum path: The path of the object. etl_args: Per-request transform arguments to forward to the next stage. Returns: - status code of the direct put request, transformed data, length of the transformed data (if any) + status code of the direct put request, transformed data, length of the + transformed data (if any), and whether the ack carried + HEADER_DIRECT_PUT_COMPLETE """ try: url = compose_etl_direct_put_url( @@ -207,7 +216,7 @@ def _direct_put( # pylint: disable=too-many-arguments,too-many-positional-argum self.server.etl_server.logger.error( "Exception during direct put to %s: %s", direct_put_url, e ) - return STATUS_INTERNAL_SERVER_ERROR, error, 0 + return STATUS_INTERNAL_SERVER_ERROR, error, 0, False def _get_fqn_content(self, path: str) -> bytes: """ @@ -249,7 +258,7 @@ def _direct_put_with_retry( # pylint: disable=too-many-arguments,too-many-posit remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """Buffered direct-put with exponential-backoff retry on transient errors.""" etl = self.server.etl_server for attempt in range(etl.direct_put_retries + 1): @@ -280,7 +289,7 @@ def _direct_put_stream_with_retry( # pylint: disable=too-many-arguments,too-man etl_args: str, is_get: bool, remaining_pipeline: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """ Streaming direct-put with exponential-backoff retry on transient errors. @@ -352,7 +361,7 @@ def _direct_put_stream( # pylint: disable=too-many-arguments,too-many-positiona remaining_pipeline: str = "", path: str = "", etl_args: str = "", - ) -> Tuple[int, bytes, int]: + ) -> Tuple[int, bytes, int, bool]: """Stream transformed output directly to the next pipeline stage.""" try: url = compose_etl_direct_put_url( @@ -384,7 +393,7 @@ def _direct_put_stream( # pylint: disable=too-many-arguments,too-many-positiona root, exc_info=True, ) - return STATUS_INTERNAL_SERVER_ERROR, str(e).encode(), 0 + return STATUS_INTERNAL_SERVER_ERROR, str(e).encode(), 0, False def _send_with_pipeline( self, transformed: bytes, path: str, etl_args: str = "" @@ -402,7 +411,7 @@ def _send_with_pipeline( if pipeline_header: first_url, remaining_pipeline = parse_etl_pipeline(pipeline_header) if first_url: - status_code, transformed, direct_put_length = ( + status_code, transformed, direct_put_length, direct_put_complete = ( self._direct_put_with_retry( first_url, transformed, remaining_pipeline, path, etl_args ) @@ -411,6 +420,7 @@ def _send_with_pipeline( status_code=status_code, length=len(transformed), direct_put_length=direct_put_length, + direct_put_complete=direct_put_complete, ) if transformed: self.wfile.write(transformed) @@ -436,6 +446,7 @@ def _handle_streaming(self, fqn, raw_path, etl_args, is_get): status_code=result[0], length=len(result[1]), direct_put_length=result[2], + direct_put_complete=result[3], ) if result[1]: self.wfile.write(result[1]) diff --git a/python/tests/integration/sdk/test_etl_webserver_pipeline.py b/python/tests/integration/sdk/test_etl_webserver_pipeline.py index 7e25a4130c..202a9498d5 100644 --- a/python/tests/integration/sdk/test_etl_webserver_pipeline.py +++ b/python/tests/integration/sdk/test_etl_webserver_pipeline.py @@ -13,6 +13,7 @@ from aistore.sdk.const import ( HEADER_NODE_URL, + HEADER_DIRECT_PUT_COMPLETE, HEADER_DIRECT_PUT_LENGTH, HEADER_CONTENT_LENGTH, ETL_WS_PIPELINE, @@ -93,6 +94,18 @@ def get_mime_type(self) -> str: return "application/flask" +def make_target_ack(length: int, mock_cls=Mock): + """Mock of the target's delivered ack: 204 + Ais-Direct-Put-Complete + length.""" + ack = mock_cls() + ack.status_code = 204 + ack.content = b"" + ack.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(length), + } + return ack + + class TestPipelineBase(unittest.TestCase): def setUp(self): os.environ["AIS_TARGET_URL"] = "http://localhost:8080" @@ -353,29 +366,25 @@ def test_http_to_target_pipeline_chain(self): server1 = self._start_http_server(12001, "step1") server2 = self._start_http_server(12002, "step2") - # Create mock response for target server call - target_response = Mock() - target_response.status_code = 200 # Will be converted to 204 by server - target_response.content = b"" - target_response.headers = {} + content = b"original" + result = server2.transform(server1.transform(content)) server2.client_put = Mock() - server2.client_put.return_value = target_response + server2.client_put.return_value = make_target_ack(len(result)) # Create pipeline header: server1 -> server2 -> target pipeline = "http://localhost:12002/transform,http://localhost:12003/target" headers = {HEADER_NODE_URL: pipeline} - content = b"original" - result = server2.transform(server1.transform(content)) # Send request to first server with pipeline response = requests.put( "http://localhost:12001/test", data=content, headers=headers, timeout=5 ) - # Should get 204 (server converts target's 200 with empty content to 204) + # The target's marked 204 ack is propagated back through the pipeline self.assertEqual(response.status_code, 204) self.assertEqual(response.content, b"") + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) self.assertEqual( response.headers.get(HEADER_DIRECT_PUT_LENGTH), str(len(result)) ) @@ -420,29 +429,25 @@ def test_flask_to_target_pipeline_chain(self): server1 = self._start_flask_server(19021, "step1") server2 = self._start_flask_server(19022, "step2") - # Create mock response for target server call - target_response = Mock() - target_response.status_code = 200 # Will be converted to 204 by server - target_response.content = b"" - target_response.headers = {} + content = b"original" + result = server2.transform(server1.transform(content, "", ""), "", "") server2.client_put = Mock() - server2.client_put.return_value = target_response + server2.client_put.return_value = make_target_ack(len(result)) # Create pipeline header: server1 -> server2 -> target pipeline = "http://localhost:19022/transform,http://localhost:19023/target" headers = {HEADER_NODE_URL: pipeline} - content = b"original" - result = server2.transform(server1.transform(content, "", ""), "", "") # Send request to first server with pipeline response = requests.put( "http://localhost:19021/test", data=content, headers=headers, timeout=5 ) - # Should get 204 (server converts target's 200 with empty content to 204) + # The target's marked 204 ack is propagated back through the pipeline self.assertEqual(response.status_code, 204) self.assertEqual(response.content, b"") + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) self.assertEqual( response.headers.get(HEADER_DIRECT_PUT_LENGTH), str(len(result)) ) @@ -487,30 +492,25 @@ def test_fastapi_to_target_pipeline_chain(self): server1 = self._start_fastapi_server(19041, "step1") server2 = self._start_fastapi_server(19042, "step2") - # Create mock response for target server call - target_response = AsyncMock() - target_response.status_code = 200 # Will be converted to 204 by server - target_response.content = b"" - target_response.headers = {} + content = b"original" + result = server2.transform(server1.transform(content, "", ""), "", "") - # Mock the direct delivery response (simulate 200 OK) server2.client.put = AsyncMock() - server2.client.put.return_value = target_response + server2.client.put.return_value = make_target_ack(len(result), AsyncMock) # Create pipeline header: server1 -> server2 -> target pipeline = "http://localhost:19042/transform,http://localhost:19043/target" headers = {HEADER_NODE_URL: pipeline} - content = b"original" - result = server2.transform(server1.transform(content, "", ""), "", "") # Send request to first server with pipeline response = requests.put( "http://localhost:19041/test", data=content, headers=headers, timeout=5 ) - # Should get 204 (server converts target's 200 with empty content to 204) + # The target's marked 204 ack is propagated back through the pipeline self.assertEqual(response.status_code, 204) self.assertEqual(response.content, b"") + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) self.assertEqual( response.headers.get(HEADER_DIRECT_PUT_LENGTH), str(len(result)) ) @@ -521,6 +521,95 @@ def test_fastapi_to_target_pipeline_chain(self): headers={HEADER_CONTENT_LENGTH: str(len(result))}, ) + @pytest.mark.etl + def test_http_to_legacy_target_pipeline_chain(self): + """A legacy target's bare 200 + Content-Length: 0 ack still reads as + delivered (fallback path), and the outgoing 204 gains NO marker.""" + + server1 = self._start_http_server(12011, "step1") + server2 = self._start_http_server(12012, "step2") + + content = b"original" + result = server2.transform(server1.transform(content)) + + # Legacy target ack: bare 200 with Content-Length: 0, no marker + target_response = Mock() + target_response.status_code = 200 + target_response.content = b"" + target_response.headers = {HEADER_CONTENT_LENGTH: "0"} + + server2.client_put = Mock() + server2.client_put.return_value = target_response + + pipeline = "http://localhost:12012/transform,http://localhost:12013/target" + headers = {HEADER_NODE_URL: pipeline} + + response = requests.put( + "http://localhost:12011/test", data=content, headers=headers, timeout=5 + ) + + self.assertEqual(response.status_code, 204) + self.assertEqual(response.content, b"") + self.assertNotIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) + self.assertEqual( + response.headers.get(HEADER_DIRECT_PUT_LENGTH), str(len(result)) + ) + + @pytest.mark.etl + def test_multi_stage_marked_ack_propagation(self): + """The target's marked 204 ack survives two intermediate ETL stages.""" + + server1 = self._start_http_server(12021, "step1") + server2 = self._start_http_server(12022, "step2") + server3 = self._start_http_server(12023, "step3") + + content = b"original" + result = server3.transform(server2.transform(server1.transform(content))) + + server3.client_put = Mock() + server3.client_put.return_value = make_target_ack(len(result)) + + pipeline = ( + "http://localhost:12022/transform," + "http://localhost:12023/transform," + "http://localhost:12024/target" + ) + headers = {HEADER_NODE_URL: pipeline} + + response = requests.put( + "http://localhost:12021/test", data=content, headers=headers, timeout=5 + ) + + self.assertEqual(response.status_code, 204) + self.assertEqual(response.content, b"") + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) + self.assertEqual( + response.headers.get(HEADER_DIRECT_PUT_LENGTH), str(len(result)) + ) + + @pytest.mark.etl + def test_target_zero_length_ack_propagation(self): + """A marked ack for an empty stored object propagates + Ais-Direct-Put-Length: 0 instead of dropping the header.""" + + self._start_http_server(12031, "step1") + server2 = self._start_http_server(12032, "step2") + + server2.client_put = Mock() + server2.client_put.return_value = make_target_ack(0) + + pipeline = "http://localhost:12032/transform,http://localhost:12033/target" + headers = {HEADER_NODE_URL: pipeline} + + response = requests.put( + "http://localhost:12031/test", data=b"original", headers=headers, timeout=5 + ) + + self.assertEqual(response.status_code, 204) + self.assertEqual(response.content, b"") + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) + self.assertEqual(response.headers.get(HEADER_DIRECT_PUT_LENGTH), "0") + @pytest.mark.etl def test_mixed_server_type_pipeline(self): """Test pipeline forwarding across different server types.""" @@ -776,16 +865,15 @@ def test_websocket_to_target_direct_pipeline(self): fastapi_server = self._start_fastapi_server(14031, "ws_target") client = TestClient(fastapi_server.app) - # Mock target response + test_data = b"websocket_target_data" + expected_len = len(fastapi_server.transform(test_data, "", "")) + + # Mock the target's delivered ack (204 + Ais-Direct-Put-Complete) with patch.object(fastapi_server, "client", new=AsyncMock()) as mock_client: - mock_resp = AsyncMock() - mock_resp.status_code = 200 - mock_resp.content = b"" - mock_client.put.return_value = mock_resp + mock_client.put.return_value = make_target_ack(expected_len, AsyncMock) # Test direct pipeline to target with client.websocket_connect("/ws") as websocket: - test_data = b"websocket_target_data" pipeline = "http://localhost:14032/target" websocket.send_json(data={ETL_WS_PIPELINE: pipeline}, mode="binary") diff --git a/python/tests/unit/sdk/test_etl_webserver.py b/python/tests/unit/sdk/test_etl_webserver.py index c9d8a110a7..f44f57e246 100644 --- a/python/tests/unit/sdk/test_etl_webserver.py +++ b/python/tests/unit/sdk/test_etl_webserver.py @@ -32,6 +32,7 @@ HEADER_CONTENT_LENGTH, ETL_WS_FQN, ETL_WS_PIPELINE, + HEADER_DIRECT_PUT_COMPLETE, HEADER_DIRECT_PUT_LENGTH, HEADER_ETL_RETRY_REASON, ETL_RETRY_REASON_DIRECT_PUT_TRANSIENT, @@ -190,16 +191,21 @@ def test_transform_get_with_direct_put(self): mock_get_resp.content = b"original" handler.server.etl_server.session.get.return_value = mock_get_resp - # Simulate direct put success (200) + # Simulate direct put success (204 + Ais-Direct-Put-Complete = delivered) mock_put_resp = MagicMock() - mock_put_resp.status_code = 200 + mock_put_resp.status_code = 204 mock_put_resp.content = b"" + mock_put_resp.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(len(b"transformed")), + } handler.server.etl_server.client_put.return_value = mock_put_resp handler.do_GET() handler.server.etl_server.client_put.assert_called_with( direct_put_url, b"transformed", headers={} ) handler.send_response.assert_called_with(204) + handler.send_header.assert_any_call(HEADER_DIRECT_PUT_COMPLETE, "true") handler.send_header.assert_called_with( HEADER_DIRECT_PUT_LENGTH, str(len(b"transformed")) ) @@ -222,16 +228,21 @@ def test_transform_put_with_direct_put(self): handler = DummyRequestHandler() handler.headers = {HEADER_NODE_URL: direct_put_url} - # Simulate direct put success (200) + # Simulate direct put success (204 + Ais-Direct-Put-Complete = delivered) mock_put_resp = MagicMock() - mock_put_resp.status_code = 200 + mock_put_resp.status_code = 204 mock_put_resp.content = b"" + mock_put_resp.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(len(b"transformed")), + } handler.server.etl_server.client_put.return_value = mock_put_resp handler.do_PUT() handler.server.etl_server.client_put.assert_called_with( direct_put_url, b"transformed", headers={} ) handler.send_response.assert_called_with(204) + handler.send_header.assert_any_call(HEADER_DIRECT_PUT_COMPLETE, "true") handler.send_header.assert_called_with( HEADER_DIRECT_PUT_LENGTH, str(len(b"transformed")) ) @@ -401,10 +412,14 @@ def test_hpush_with_direct_put(self): input_content = b"input data" transformed_content = self.etl_server.transform(input_content, path, "") - # Mock the direct delivery response (simulate 200 OK) + # Mock the direct delivery ack (204 + Ais-Direct-Put-Complete = delivered) mock_response_success = AsyncMock() mock_response_success.content = b"" - mock_response_success.status_code = 200 + mock_response_success.status_code = 204 + mock_response_success.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(len(transformed_content)), + } self.etl_server.client = AsyncMock() self.etl_server.client.put.return_value = mock_response_success @@ -413,6 +428,7 @@ def test_hpush_with_direct_put(self): self.assertEqual(response.status_code, 204) self.assertEqual(response.content, b"") # No content returned + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) # propagated self.assertEqual( response.headers.get(HEADER_DIRECT_PUT_LENGTH), str(len(transformed_content)), @@ -433,6 +449,29 @@ def test_hpush_with_direct_put(self): self.assertEqual(response.content, b"error message") self.etl_server.client.put.assert_awaited_once() + @unittest.skipIf(sys.version_info < (3, 9), "requires Python 3.9 or higher") + def test_hpush_with_direct_put_empty_object(self): + """A zero-length ack (empty object stored) propagates the marker and + Ais-Direct-Put-Length: 0 instead of dropping the header.""" + path = "test/object" + + mock_response_success = AsyncMock() + mock_response_success.content = b"" + mock_response_success.status_code = 204 + mock_response_success.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: "0", + } + self.etl_server.client = AsyncMock() + self.etl_server.client.put.return_value = mock_response_success + + headers = {HEADER_NODE_URL: "http://localhost:8080/ais/@/etl_dst/test/object"} + response = self.client.put(f"/{path}", content=b"input data", headers=headers) + + self.assertEqual(response.status_code, 204) + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) + self.assertEqual(response.headers.get(HEADER_DIRECT_PUT_LENGTH), "0") + @unittest.skipIf(sys.version_info < (3, 9), "requires Python 3.9 or higher") def test_hpush_with_direct_put_and_fqn(self): path = "test/object" @@ -440,7 +479,7 @@ def test_hpush_with_direct_put_and_fqn(self): input_content = b"input data" transformed_content = self.etl_server.transform(input_content, path, "") - # Mock the direct put response (simulate 200 OK) + # Mock the direct put ack (204 + Ais-Direct-Put-Complete = delivered) with patch.object( self.etl_server, "_get_fqn_content", @@ -448,7 +487,11 @@ def test_hpush_with_direct_put_and_fqn(self): ) as get_fqn_mock: mock_response_success = AsyncMock() mock_response_success.content = b"" - mock_response_success.status_code = 200 + mock_response_success.status_code = 204 + mock_response_success.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(len(transformed_content)), + } self.etl_server.client = AsyncMock() self.etl_server.client.put.return_value = mock_response_success @@ -462,6 +505,7 @@ def test_hpush_with_direct_put_and_fqn(self): self.assertEqual(response.status_code, 204) self.assertEqual(response.content, b"") # No content returned + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) # propagated self.assertEqual( response.headers.get(HEADER_DIRECT_PUT_LENGTH), str(len(transformed_content)), @@ -499,11 +543,15 @@ def test_websocket_with_direct_put(self): input_data = b"testdata" direct_put_url = "http://localhost:8080/ais/@/etl_dst/final" - # Mock the direct put response (simulate 200 OK) => return length as ACK + # Mock the direct put ack (204 + marker) => return length as ACK with patch.object(self.etl_server, "client", new=AsyncMock()) as mock_client: mock_resp = AsyncMock() - mock_resp.status_code = 200 + mock_resp.status_code = 204 mock_resp.content = b"" + mock_resp.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(len(input_data)), + } mock_client.put.return_value = mock_resp with self.client.websocket_connect("/ws") as websocket: @@ -568,11 +616,15 @@ def test_websocket_with_direct_put_and_fqn(self): original_content = b"original data" direct_put_url = "http://localhost:8080/ais/@/etl_dst/final" transformed_content = self.etl_server.transform(original_content, fqn, "") - # Mock the direct put response (simulate 200 OK) + # Mock the direct put ack (204 + Ais-Direct-Put-Complete = delivered) with patch.object(self.etl_server, "client", new=AsyncMock()) as mock_client: mock_resp = AsyncMock() - mock_resp.status_code = 200 + mock_resp.status_code = 204 mock_resp.content = b"" + mock_resp.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(len(transformed_content)), + } mock_client.put.return_value = mock_resp with patch.object( @@ -702,11 +754,19 @@ def test_direct_put_delivery(self): headers = {HEADER_NODE_URL: "http://localhost:8080/ais/@/etl_dst/test/object"} with patch("requests.Session.put") as mock_put: - # Mock the direct delivery response (simulate 200 OK) - mock_put.return_value = MagicMock(status_code=200, content=b"") + # Mock the direct delivery ack (204 + Ais-Direct-Put-Complete) + mock_put.return_value = MagicMock( + status_code=204, + content=b"", + headers={ + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(len(transformed_content)), + }, + ) response = self.client.put(f"/{path}", data=input_content, headers=headers) self.assertEqual(response.status_code, 204) + self.assertIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) # propagated self.assertEqual( response.headers.get(HEADER_DIRECT_PUT_LENGTH), str(len(transformed_content)), @@ -721,6 +781,30 @@ def test_direct_put_delivery(self): self.assertEqual(response.status_code, 500) self.assertEqual(response.data, b"error message") + @unittest.skipIf(sys.version_info < (3, 9), "requires Python 3.9 or higher") + def test_direct_put_delivery_legacy_target(self): + """A legacy target's bare 200 + Content-Length: 0 ack is still read as + delivered, and the outgoing 204 must NOT gain the marker (the marker + is only echoed when received).""" + path = "test/object" + input_content = b"input data" + transformed_content = self.etl_server.transform(input_content, path, "") + headers = {HEADER_NODE_URL: "http://localhost:8080/ais/@/etl_dst/test/object"} + + with patch("requests.Session.put") as mock_put: + mock_put.return_value = MagicMock( + status_code=200, content=b"", headers={HEADER_CONTENT_LENGTH: "0"} + ) + response = self.client.put(f"/{path}", data=input_content, headers=headers) + + self.assertEqual(response.status_code, 204) + self.assertNotIn(HEADER_DIRECT_PUT_COMPLETE, response.headers) + self.assertEqual( + response.headers.get(HEADER_DIRECT_PUT_LENGTH), + str(len(transformed_content)), + ) + self.assertEqual(response.data, b"") + class TestBaseEnforcement(unittest.TestCase): def test_fastapi_server_without_target_url(self): @@ -1339,8 +1423,12 @@ async def fake_put(_url, *, content=None, headers=None): async for chunk in content: received.extend(chunk) mock_resp = MagicMock() - mock_resp.status_code = 200 + mock_resp.status_code = 204 mock_resp.content = b"" + mock_resp.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: str(len(received)), + } return mock_resp self.etl_server.client = AsyncMock() @@ -1892,10 +1980,14 @@ async def _stream(): return req def _ok_response(self) -> MagicMock: + # target's delivered ack: 204 + Ais-Direct-Put-Complete + length resp = MagicMock() - resp.status_code = 200 + resp.status_code = 204 resp.content = b"" - resp.headers = {} + resp.headers = { + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: "4", + } return resp async def _call(self, req, retries=None, is_get=True, fqn=""): @@ -1925,7 +2017,7 @@ async def test_succeeds_on_first_attempt(self): self.server.client.put.return_value = self._ok_response() with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: result = await self._call(self._make_request()) - # 200 + empty content → handle_direct_put_response returns 204 + # marked ack → handle_direct_put_response returns 204 self.assertEqual(result[0], 204) self.server.client.put.assert_awaited_once() mock_sleep.assert_not_called() @@ -1945,7 +2037,7 @@ async def mock_put(*_, **__): self.server.client.put.side_effect = mock_put with patch("asyncio.sleep", new_callable=AsyncMock): result = await self._call(self._make_request()) - self.assertEqual(result[0], 204) # 200 + empty content → 204 + self.assertEqual(result[0], 204) # marked ack → 204 self.assertEqual(call_count, 3) @unittest.skipIf(sys.version_info < (3, 9), "requires Python 3.9 or higher") @@ -2158,16 +2250,18 @@ def _put(self, url="http://target/obj", data=b"data"): def test_succeeds_on_first_attempt(self): """No retry when first attempt succeeds.""" - with patch.object(self.server, "_direct_put", return_value=(200, b"ok", 2)): + with patch.object( + self.server, "_direct_put", return_value=(200, b"ok", 2, False) + ): with patch("time.sleep") as mock_sleep: result = self._retry() - self.assertEqual(result, (200, b"ok", 2)) + self.assertEqual(result, (200, b"ok", 2, False)) mock_sleep.assert_not_called() def test_retries_on_transient_error_then_succeeds(self): """Retries on ETLDirectPutTransientError and succeeds on second attempt.""" self.server.direct_put_retries = 2 - ok = (200, b"", 4) + ok = (200, b"", 4, False) side_effects = [ ETLDirectPutTransientError("http://target/obj", requests.ConnectionError()), ok, @@ -2208,7 +2302,9 @@ def test_direct_put_wraps_chunked_encoding_error(self): def test_non_transient_error_returns_500(self): """Non-transient exceptions in _direct_put return 500 without retrying.""" - with patch.object(self.server, "_direct_put", return_value=(500, b"err", 0)): + with patch.object( + self.server, "_direct_put", return_value=(500, b"err", 0, False) + ): with patch("time.sleep") as mock_sleep: result = self._retry() self.assertEqual(result[0], 500) @@ -2224,7 +2320,7 @@ def side_effect(*_args, **_kwargs): call_count += 1 if call_count < 4: raise ETLDirectPutTransientError("url", requests.ConnectionError()) - return (200, b"", 0) + return (200, b"", 0, False) with patch.object(self.server, "_direct_put", side_effect=side_effect): with patch("time.sleep") as mock_sleep: @@ -2274,16 +2370,18 @@ def _put(self, url="http://target/obj", data=b"data", path="/obj"): def test_succeeds_on_first_attempt(self): """No retry when first attempt succeeds.""" - with patch.object(self.handler, "_direct_put", return_value=(200, b"ok", 5)): + with patch.object( + self.handler, "_direct_put", return_value=(200, b"ok", 5, False) + ): with patch("time.sleep") as mock_sleep: result = self._retry() - self.assertEqual(result, (200, b"ok", 5)) + self.assertEqual(result, (200, b"ok", 5, False)) mock_sleep.assert_not_called() def test_retries_on_transient_error_then_succeeds(self): """Retries on ETLDirectPutTransientError and succeeds on second attempt.""" self.handler.server.etl_server.direct_put_retries = 2 - ok = (200, b"ok", 3) + ok = (200, b"ok", 3, False) side_effects = [ ETLDirectPutTransientError("url", requests.ConnectionError()), ok, @@ -2320,7 +2418,9 @@ def test_direct_put_wraps_chunked_encoding_error(self): def test_non_transient_error_returns_500(self): """Non-transient exceptions in _direct_put return 500 without retrying.""" - with patch.object(self.handler, "_direct_put", return_value=(500, b"err", 0)): + with patch.object( + self.handler, "_direct_put", return_value=(500, b"err", 0, False) + ): with patch("time.sleep") as mock_sleep: result = self._retry() self.assertEqual(result[0], 500) @@ -2336,7 +2436,7 @@ def side_effect(*_args, **_kwargs): call_count += 1 if call_count < 4: raise ETLDirectPutTransientError("url", requests.ConnectionError()) - return (200, b"", 0) + return (200, b"", 0, False) with patch.object(self.handler, "_direct_put", side_effect=side_effect): with patch("time.sleep") as mock_sleep: @@ -2398,17 +2498,17 @@ def test_succeeds_on_first_attempt(self): self.server, "_get_stream_reader", return_value=self._make_reader() ): with patch.object( - self.server, "_direct_put_stream", return_value=(204, b"", 5) + self.server, "_direct_put_stream", return_value=(204, b"", 5, True) ): with patch("time.sleep") as mock_sleep: result = self._call() - self.assertEqual(result, (204, b"", 5)) + self.assertEqual(result, (204, b"", 5, True)) mock_sleep.assert_not_called() def test_retries_on_transient_error_then_succeeds(self): """Retries on ETLDirectPutTransientError and succeeds on second attempt.""" self.server.direct_put_retries = 2 - ok = (204, b"", 4) + ok = (204, b"", 4, True) call_count = 0 def put_side(*_args, **_kwargs): @@ -2462,7 +2562,7 @@ def put_side(*_args, **_kwargs): raise ETLDirectPutTransientError( self._DIRECT_PUT_URL, requests.ConnectionError() ) - return (204, b"", 0) + return (204, b"", 0, True) readers = [self._make_reader() for _ in range(3)] reader_iter = iter(readers) @@ -2482,7 +2582,7 @@ def test_reader_always_closed_on_success(self): mock_reader = self._make_reader() with patch.object(self.server, "_get_stream_reader", return_value=mock_reader): with patch.object( - self.server, "_direct_put_stream", return_value=(204, b"", 0) + self.server, "_direct_put_stream", return_value=(204, b"", 0, True) ): with patch.object( self.server, "close_reader", wraps=self.server.close_reader @@ -2530,7 +2630,7 @@ def put_side(*_args, **_kwargs): raise ETLDirectPutTransientError( self._DIRECT_PUT_URL, requests.ConnectionError() ) - return (204, b"", 0) + return (204, b"", 0, True) readers = [self._make_reader() for _ in range(4)] reader_iter = iter(readers) @@ -2655,17 +2755,17 @@ def test_succeeds_on_first_attempt(self): self.handler, "_get_stream_reader", return_value=self._make_reader() ): with patch.object( - self.handler, "_direct_put_stream", return_value=(204, b"", 5) + self.handler, "_direct_put_stream", return_value=(204, b"", 5, True) ): with patch("time.sleep") as mock_sleep: result = self._call() - self.assertEqual(result, (204, b"", 5)) + self.assertEqual(result, (204, b"", 5, True)) mock_sleep.assert_not_called() def test_retries_on_transient_error_then_succeeds(self): """Replayable GET path retries on transient error and succeeds on second attempt.""" self.handler.server.etl_server.direct_put_retries = 2 - ok = (204, b"", 4) + ok = (204, b"", 4, True) call_count = 0 def put_side(*_args, **_kwargs): @@ -2749,7 +2849,7 @@ def test_reader_always_closed_on_success(self): mock_reader = self._make_reader() with patch.object(self.handler, "_get_stream_reader", return_value=mock_reader): with patch.object( - self.handler, "_direct_put_stream", return_value=(204, b"", 0) + self.handler, "_direct_put_stream", return_value=(204, b"", 0, True) ): self._call() self.handler.server.etl_server.close_reader.assert_called_once_with(mock_reader) @@ -2785,7 +2885,7 @@ def put_side(*_args, **_kwargs): raise ETLDirectPutTransientError( self._DIRECT_PUT_URL, requests.ConnectionError() ) - return (204, b"", 0) + return (204, b"", 0, True) readers = [self._make_reader() for _ in range(4)] reader_iter = iter(readers) @@ -2848,7 +2948,7 @@ def test_replayable_exhausted_does_not_set_bail_flag(self): def test_fqn_put_still_retries(self): """FQN-backed PUT is replayable; retries proceed normally.""" self.handler.server.etl_server.direct_put_retries = 2 - ok = (204, b"", 4) + ok = (204, b"", 4, True) call_count = 0 def put_side(*_args, **_kwargs): @@ -2877,7 +2977,7 @@ def put_side(*_args, **_kwargs): def test_get_path_still_retries(self): """GET path stays replayable; retries proceed normally (regression guard).""" self.handler.server.etl_server.direct_put_retries = 2 - ok = (204, b"", 4) + ok = (204, b"", 4, True) call_count = 0 def put_side(*_args, **_kwargs): @@ -3060,7 +3160,7 @@ def test_connection_refused_returns_502(self): """_direct_put() returns (502, ...) for ConnectionRefused — not ETLDirectPutTransientError.""" exc = _make_connection_refused_error() with patch.object(self.server, "client_put", side_effect=exc): - status, body, _ = self._put() + status, body, _, _ = self._put() self.assertEqual(status, 502) self.assertIn(b"ConnectionError", body) self.assertIn(b"/nonexistent", body) @@ -3070,7 +3170,7 @@ def test_connection_refused_not_retried(self): exc = _make_connection_refused_error() with patch.object(self.server, "client_put", side_effect=exc): with patch("time.sleep") as mock_sleep: - status, _, _ = self._retry() + status, _, _, _ = self._retry() self.assertEqual(status, 502) mock_sleep.assert_not_called() @@ -3132,7 +3232,7 @@ def test_connection_refused_returns_502(self): """_direct_put() returns (502, ...) for ConnectionRefused — not ETLDirectPutTransientError.""" exc = _make_connection_refused_error() self.handler.server.etl_server.client_put = MagicMock(side_effect=exc) - status, body, _ = self._put() + status, body, _, _ = self._put() self.assertEqual(status, 502) self.assertIn(b"ConnectionError", body) self.assertIn(b"/nonexistent", body) @@ -3142,7 +3242,7 @@ def test_connection_refused_not_retried(self): exc = _make_connection_refused_error() self.handler.server.etl_server.client_put = MagicMock(side_effect=exc) with patch("time.sleep") as mock_sleep: - status, _, _ = self._retry() + status, _, _, _ = self._retry() self.assertEqual(status, 502) mock_sleep.assert_not_called() @@ -3216,7 +3316,14 @@ def test_http_buffered_forwards_etl_args(self): handler.path = "/test/object?etl_args=jpeg" handler.headers = {HEADER_NODE_URL: "http://some-target/put/object"} - mock_put_resp = MagicMock(status_code=200, content=b"") + mock_put_resp = MagicMock( + status_code=204, + content=b"", + headers={ + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: "10", + }, + ) handler.server.etl_server.client_put.return_value = mock_put_resp handler.do_PUT() @@ -3232,7 +3339,14 @@ def test_fastapi_buffered_forwards_etl_args(self): server = DummyFastAPIServer() client = TestClient(server.app) server.client = AsyncMock() - server.client.put.return_value = AsyncMock(status_code=200, content=b"") + server.client.put.return_value = AsyncMock( + status_code=204, + content=b"", + headers={ + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: "10", + }, + ) headers = {HEADER_NODE_URL: "http://localhost:8080/ais/@/etl_dst/test/object"} response = client.put( @@ -3253,7 +3367,14 @@ def test_fastapi_websocket_forwards_etl_args(self): direct_put_url = "http://localhost:8080/ais/@/etl_dst/final" with patch.object(server, "client", new=AsyncMock()) as mock_client: - mock_client.put.return_value = AsyncMock(status_code=200, content=b"") + mock_client.put.return_value = AsyncMock( + status_code=204, + content=b"", + headers={ + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: "8", + }, + ) with client.websocket_connect("/ws") as websocket: websocket.send_json( data={ETL_WS_PIPELINE: direct_put_url, QPARAM_ETL_ARGS: "jpeg"}, @@ -3275,7 +3396,14 @@ def test_flask_buffered_forwards_etl_args(self): headers = {HEADER_NODE_URL: "http://localhost:8080/ais/@/etl_dst/test/object"} with patch("requests.Session.put") as mock_put: - mock_put.return_value = MagicMock(status_code=200, content=b"") + mock_put.return_value = MagicMock( + status_code=204, + content=b"", + headers={ + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: "10", + }, + ) response = client.put( "/test/object?etl_args=jpeg", data=b"input data", headers=headers ) @@ -3290,7 +3418,14 @@ def test_flask_streaming_forwards_etl_args(self): """FlaskServer forwards etl_args on the streaming direct-put hop.""" server = DummyFlaskServer() server.session = MagicMock() - server.session.put.return_value = MagicMock(status_code=200, content=b"") + server.session.put.return_value = MagicMock( + status_code=204, + content=b"", + headers={ + HEADER_DIRECT_PUT_COMPLETE: "true", + HEADER_DIRECT_PUT_LENGTH: "5", + }, + ) # pylint: disable=protected-access server._direct_put_stream( @@ -3305,3 +3440,145 @@ def test_flask_streaming_forwards_etl_args(self): self.assertEqual( parse_qs(urlparse(called_url).query)[QPARAM_ETL_ARGS], ["jpeg"] ) + + +class TestHandleDirectPutResponse(unittest.TestCase): + """Direct unit coverage of `ETLServer.handle_direct_put_response`. + + Classification is marker-first: a response carrying + `HEADER_DIRECT_PUT_COMPLETE` (presence-based) is the target's delivered + ack and is propagated as 204 + marker + `HEADER_DIRECT_PUT_LENGTH`. + Markerless responses fall through to the legacy handling, unchanged: + the 200 classification keyed on `Content-Length` (`0` means delivered; + absent/non-zero means content, forwarded as-is), kept for targets that + predate the marker. + """ + + def setUp(self): + env = mock.patch.dict(os.environ, {"AIS_TARGET_URL": "http://localhost:8080"}) + env.start() + self.addCleanup(env.stop) + self.server = DummyFastAPIServer() + + @staticmethod + def _resp(status_code, content=b"", headers=None): + resp = MagicMock() + resp.status_code = status_code + resp.content = content + resp.headers = {} if headers is None else headers + return resp + + def test_204_with_marker_is_delivered(self): + resp = self._resp( + 204, + b"", + {HEADER_DIRECT_PUT_COMPLETE: "true", HEADER_DIRECT_PUT_LENGTH: "1234"}, + ) + self.assertEqual( + self.server.handle_direct_put_response(resp, b""), + (204, b"", 1234, True), + ) + + def test_marker_with_length_zero_is_delivered(self): + # empty object stored: the zero length is propagated, not dropped + resp = self._resp( + 204, + b"", + {HEADER_DIRECT_PUT_COMPLETE: "true", HEADER_DIRECT_PUT_LENGTH: "0"}, + ) + self.assertEqual( + self.server.handle_direct_put_response(resp, b""), + (204, b"", 0, True), + ) + + def test_marker_with_empty_value_is_delivered(self): + # presence-based: an empty header value still counts + resp = self._resp( + 204, + b"", + {HEADER_DIRECT_PUT_COMPLETE: "", HEADER_DIRECT_PUT_LENGTH: "7"}, + ) + self.assertEqual( + self.server.handle_direct_put_response(resp, b""), + (204, b"", 7, True), + ) + + def test_marker_missing_length_defaults_to_zero(self): + # matches the legacy 204 branch and Go's directPut: absent length + # header degrades to 0 rather than failing a stored object + resp = self._resp(204, b"", {HEADER_DIRECT_PUT_COMPLETE: "true"}) + self.assertEqual( + self.server.handle_direct_put_response(resp, b""), + (204, b"", 0, True), + ) + + def test_marker_on_200_is_normalized_to_204(self): + # defensive: the marker decides regardless of status + resp = self._resp( + 200, + b"", + {HEADER_DIRECT_PUT_COMPLETE: "true", HEADER_DIRECT_PUT_LENGTH: "5"}, + ) + self.assertEqual( + self.server.handle_direct_put_response(resp, b""), + (204, b"", 5, True), + ) + + def test_legacy_200_content_length_zero_is_delivered(self): + resp = self._resp(200, b"", {HEADER_CONTENT_LENGTH: "0"}) + self.assertEqual( + self.server.handle_direct_put_response(resp, b"payload"), + (204, b"", len(b"payload"), False), + ) + + def test_legacy_200_content_length_zero_uses_data_length_override(self): + resp = self._resp(200, b"", {HEADER_CONTENT_LENGTH: "0"}) + self.assertEqual( + self.server.handle_direct_put_response(resp, b"", data_length=42), + (204, b"", 42, False), + ) + + def test_legacy_200_chunked_empty_body_is_forwarded(self): + # No Content-Length (chunked): an empty body is a valid empty + # transform result and must be forwarded, not misreported as + # delivered with the original size. + resp = self._resp(200, b"") + self.assertEqual( + self.server.handle_direct_put_response(resp, b"payload"), + (200, b"", 0, False), + ) + + def test_legacy_200_chunked_body_is_forwarded(self): + resp = self._resp(200, b"transformed") + self.assertEqual( + self.server.handle_direct_put_response(resp, b"payload"), + (200, b"transformed", 0, False), + ) + + def test_legacy_200_with_content_length_is_forwarded(self): + resp = self._resp(200, b"transformed", {HEADER_CONTENT_LENGTH: "11"}) + self.assertEqual( + self.server.handle_direct_put_response(resp, b"payload"), + (200, b"transformed", 0, False), + ) + + def test_legacy_200_malformed_content_length_is_forwarded(self): + resp = self._resp(200, b"x", {HEADER_CONTENT_LENGTH: "abc"}) + self.assertEqual( + self.server.handle_direct_put_response(resp, b"payload"), + (200, b"x", 0, False), + ) + + def test_legacy_204_forwards_direct_put_length(self): + resp = self._resp(204, b"", {HEADER_DIRECT_PUT_LENGTH: "1234"}) + self.assertEqual( + self.server.handle_direct_put_response(resp, b""), + (204, b"", 1234, False), + ) + + def test_legacy_204_missing_direct_put_length_defaults_to_zero(self): + resp = self._resp(204, b"") + self.assertEqual( + self.server.handle_direct_put_response(resp, b""), + (204, b"", 0, False), + )