diff --git a/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index a56dea3cd..42c264341 100644 --- a/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -334,6 +334,9 @@ async def _handle_execution( except SmithyRetryException: raise context_with_response.response await sleep(retry_token.retry_delay) + current_body = context_with_transport_request.transport_request.body + if (seek := getattr(current_body, "seek", None)) is not None: + await seek(0) else: # Step 8: Invoke record_success retry_strategy.record_success(token=retry_token) diff --git a/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/integration/RestJsonProtocolGenerator.java b/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/integration/RestJsonProtocolGenerator.java index e916980da..066c7a76e 100644 --- a/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/integration/RestJsonProtocolGenerator.java +++ b/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/integration/RestJsonProtocolGenerator.java @@ -138,7 +138,7 @@ protected void serializeDocumentBody( ) { writer.addDependency(SmithyPythonDependency.SMITHY_JSON); writer.addImport("smithy_json", "JSONCodec"); - writer.addImport("smithy_core.aio.types", "AsyncBytesReader"); + writer.addImport("smithy_core.aio.types", "SeekableAsyncBytesReader"); writer.pushState(); @@ -146,7 +146,7 @@ protected void serializeDocumentBody( var body = shouldWriteDefaultBody(context, operation) ? "b'{}'" : "b''"; writer.write(""" content_length = 0 - body = AsyncBytesReader($L) + body = SeekableAsyncBytesReader($L) """, body); } else { writer.addImport("smithy_core.types", "TimestampFormat"); @@ -159,7 +159,7 @@ protected void serializeDocumentBody( content = b'{}' ${/writeDefaultBody} content_length = len(content) - body = AsyncBytesReader(content) + body = SeekableAsyncBytesReader(content) """); } @@ -223,13 +223,13 @@ protected void serializePayloadBody( return; } - writer.addImport("smithy_core.aio.types", "AsyncBytesReader"); + writer.addImport("smithy_core.aio.types", "SeekableAsyncBytesReader"); writer.write("content_length: int = 0"); CodegenUtils.accessStructureMember(context, writer, "input", payloadBinding.getMember(), () -> { if (target.isBlobShape()) { writer.write("content_length = len($L)", dataSource); - writer.write("body = AsyncBytesReader($L)", dataSource); + writer.write("body = SeekableAsyncBytesReader($L)", dataSource); return; } @@ -244,13 +244,13 @@ protected void serializePayloadBody( """, dataSource); } writer.write("content_length = len(content)"); - writer.write("body = AsyncBytesReader(content)"); + writer.write("body = SeekableAsyncBytesReader(content)"); }); if (target.isStructureShape()) { writer.write(""" else: content_length = 2 - body = AsyncBytesReader(b'{}') + body = SeekableAsyncBytesReader(b'{}') """); } }