From a03e4767d83983de23975ce8f3cd881b8a6b85d8 Mon Sep 17 00:00:00 2001 From: JordonPhillips Date: Tue, 7 Jan 2025 16:05:16 +0100 Subject: [PATCH 1/3] Use seekable bodies for requests This updates requests to use seekable bodies wherever possible so that they can be retried. --- .../integration/RestJsonProtocolGenerator.java | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) 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'{}') """); } } From 3cb180c269b833a635ade37b56b7e4d01180c161 Mon Sep 17 00:00:00 2001 From: JordonPhillips Date: Tue, 7 Jan 2025 16:28:36 +0100 Subject: [PATCH 2/3] Rewind requests before retrying them --- .../software/amazon/smithy/python/codegen/ClientGenerator.java | 2 ++ 1 file changed, 2 insertions(+) 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..0146900f9 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,8 @@ async def _handle_execution( except SmithyRetryException: raise context_with_response.response await sleep(retry_token.retry_delay) + if (seek := getattr(context_with_transport_request.transport_request.body, "seek")) is not None: + await seek(0) else: # Step 8: Invoke record_success retry_strategy.record_success(token=retry_token) From 9d6e16e24ba52d78d5bcf966eaf2db153c8acc68 Mon Sep 17 00:00:00 2001 From: Jordon Phillips Date: Wed, 8 Jan 2025 14:48:31 +0100 Subject: [PATCH 3/3] Set default in seek getattr Co-authored-by: Nate Prewitt --- .../software/amazon/smithy/python/codegen/ClientGenerator.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 0146900f9..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,7 +334,8 @@ async def _handle_execution( except SmithyRetryException: raise context_with_response.response await sleep(retry_token.retry_delay) - if (seek := getattr(context_with_transport_request.transport_request.body, "seek")) is not None: + 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