diff --git a/README.md b/README.md index a8422154d..cbda5db4c 100644 --- a/README.md +++ b/README.md @@ -161,9 +161,9 @@ from echo.models import EchoMessageInput async def main() -> None: - client = EchoService(Config(endpoint_uri="https://example.com/")) - response = await client.echo_message(EchoMessageInput(message="spam")) - print(response.message) + async with EchoService(Config(endpoint_uri="https://example.com/")) as client: + response = await client.echo_message(EchoMessageInput(message="spam")) + print(response.message) if __name__ == "__main__": diff --git a/codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java b/codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java index 929aeb602..68f5a1d91 100644 --- a/codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java +++ b/codegen/core/src/it/java/software/amazon/smithy/python/codegen/test/PythonCodegenTest.java @@ -4,6 +4,10 @@ */ package software.amazon.smithy.python.codegen.test; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import java.nio.file.Files; import java.nio.file.Path; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -20,7 +24,7 @@ public class PythonCodegenTest { @Test - public void testCodegen(@TempDir Path tempDir) { + public void testCodegen(@TempDir Path tempDir) throws IOException { // TODO: Move this to its own package once client codegen is in its own package PythonClientCodegenPlugin plugin = new PythonClientCodegenPlugin(); Model model = Model.assembler(PythonCodegenTest.class.getClassLoader()) @@ -38,5 +42,9 @@ public void testCodegen(@TempDir Path tempDir) { .model(model) .build(); plugin.execute(context); + + String client = Files.readString(tempDir.resolve("src/weather/client.py")); + assertTrue(client.contains("async def close(self) -> None:")); + assertTrue(client.contains("{id(self._config.transport): self._config.transport}")); } } diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index e9f5d7a35..993aeacb8 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java @@ -94,6 +94,27 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None): writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)), RuntimeTypes.RETRY_STRATEGY_RESOLVER); + writer.addStdlibImport("typing", "Any"); + writer.write(""" + + async def close(self) -> None: + \"\"\"Close any resources held by this client's transport.\"\"\" + await $1T(self._config.transport) + + async def __aenter__(self) -> "$2L": + return self + + async def __aexit__( + self, + exc_type: Any, + exc_value: Any, + traceback: Any, + ) -> None: + await self.close() + """, + RuntimeTypes.ASYNC_CLOSE, + serviceSymbol.getName()); + var topDownIndex = TopDownIndex.of(model); var eventStreamIndex = EventStreamIndex.of(model); for (OperationShape operation : topDownIndex.getContainedOperations(service)) { @@ -242,7 +263,10 @@ private void writeSharedOperationInit( ] if plugins: operation_plugins.extend(plugins) - config = deepcopy(self._config) + config = deepcopy( + self._config, + {id(self._config.transport): self._config.transport}, + ) for plugin in operation_plugins: plugin(config) if config.protocol is None or config.transport is None: diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java index b738779c5..c18cd679c 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/PythonSymbolProvider.java @@ -66,6 +66,7 @@ public final class PythonSymbolProvider implements SymbolProvider, ShapeVisitor< private static final Logger LOGGER = Logger.getLogger(PythonSymbolProvider.class.getName()); private static final String SHAPES_FILE = "models"; private static final String SCHEMAS_FILE = "_private/schemas"; + private static final Set CLIENT_RESERVED_METHOD_NAMES = Set.of("close"); private final Model model; private final ReservedWordSymbolProvider.Escaper escaper; @@ -297,6 +298,9 @@ public Symbol operationShape(OperationShape shape) { // Operation names are escaped like members because ultimately they're // properties on an object too. var methodName = escaper.escapeMemberName(CaseUtils.toSnakeCase(shape.getId().getName(service))); + if (CLIENT_RESERVED_METHOD_NAMES.contains(methodName)) { + methodName = escapeWord(methodName); + } var methodSymbol = createGeneratedSymbolBuilder(shape, methodName, "client", false) .putProperty(SymbolProperties.IMPORTABLE, false) .build(); diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java index 4a959e48f..7c5aed8e5 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java +++ b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java @@ -117,6 +117,7 @@ public final class RuntimeTypes { // smithy_core.aio.utils public static final Symbol ASYNC_LIST = createSymbol("aio.utils", "async_list", SmithyPythonDependency.SMITHY_CORE); + public static final Symbol ASYNC_CLOSE = createSymbol("aio.utils", "close", SmithyPythonDependency.SMITHY_CORE); // smithy_http public static final Symbol TUPLES_TO_FIELDS = diff --git a/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java b/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java index 3b8c184f2..dcf1f46e2 100644 --- a/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java +++ b/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java @@ -9,6 +9,7 @@ import org.junit.jupiter.api.Test; import software.amazon.smithy.model.Model; import software.amazon.smithy.model.shapes.MemberShape; +import software.amazon.smithy.model.shapes.OperationShape; import software.amazon.smithy.model.shapes.ShapeId; import software.amazon.smithy.model.shapes.UnionShape; @@ -84,6 +85,29 @@ public void testUnionUnknownVariantNameCollidingWithShapeUsesUnderscoreSeparator provider.toSymbol(union).expectProperty(SymbolProperties.UNION_UNKNOWN).getName()); } + @Test + public void testOperationNameCollidingWithClientMethodIsEscaped() { + Model model = loadModel(""" + $version: "2" + namespace smithy.example + + service TestService { + version: "2024-01-01" + operations: [Close] + } + + operation Close {} + """); + PythonSymbolProvider provider = createProvider(model); + var operation = model.expectShape(ShapeId.from(NS + "#Close"), OperationShape.class); + + assertEquals( + "close_", + provider.toSymbol(operation) + .expectProperty(SymbolProperties.OPERATION_METHOD) + .getName()); + } + private static Model loadModel(String smithyIdl) { return Model.assembler().addUnparsedModel("test.smithy", smithyIdl).assemble().unwrap(); } diff --git a/packages/smithy-http/.changes/next-release/smithy-http-bugfix-9b22c3201ef34610abb155ba32d0b097.json b/packages/smithy-http/.changes/next-release/smithy-http-bugfix-9b22c3201ef34610abb155ba32d0b097.json new file mode 100644 index 000000000..4e478b341 --- /dev/null +++ b/packages/smithy-http/.changes/next-release/smithy-http-bugfix-9b22c3201ef34610abb155ba32d0b097.json @@ -0,0 +1,4 @@ +{ + "type": "bugfix", + "description": "Preserved HTTP clients across operation configuration copies to avoid duplicating sessions and discarding connection pools." +} \ No newline at end of file diff --git a/packages/smithy-http/.changes/next-release/smithy-http-feature-4fe36219987843b19f163600ae3a25a2.json b/packages/smithy-http/.changes/next-release/smithy-http-feature-4fe36219987843b19f163600ae3a25a2.json new file mode 100644 index 000000000..2e1c8d504 --- /dev/null +++ b/packages/smithy-http/.changes/next-release/smithy-http-feature-4fe36219987843b19f163600ae3a25a2.json @@ -0,0 +1,4 @@ +{ + "type": "feature", + "description": "Added deterministic connection-pool cleanup and async context manager support to HTTP clients." +} diff --git a/packages/smithy-http/src/smithy_http/aio/aiohttp.py b/packages/smithy-http/src/smithy_http/aio/aiohttp.py index c5f08a1b4..37e5b13b0 100644 --- a/packages/smithy-http/src/smithy_http/aio/aiohttp.py +++ b/packages/smithy-http/src/smithy_http/aio/aiohttp.py @@ -1,8 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from copy import copy, deepcopy from itertools import chain -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Self from urllib.parse import parse_qs import yarl @@ -111,6 +110,16 @@ async def send( ) as resp: return await self._marshal_response(resp) + async def close(self) -> None: + """Close the underlying aiohttp session and its connection pool.""" + await self._session.close() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.close() + async def _prepare_body(self, body: StreamingBlob) -> AsyncBytesReader | None: """Convert a body for aiohttp, omitting seekable bodies with no data.""" if not isinstance(body, AsyncBytesReader): @@ -159,7 +168,4 @@ async def _marshal_response( ) def __deepcopy__(self, memo: Any) -> "AIOHTTPClient": - return AIOHTTPClient( - client_config=deepcopy(self._config), - _session=copy(self._session), - ) + return self diff --git a/packages/smithy-http/src/smithy_http/aio/crt.py b/packages/smithy-http/src/smithy_http/aio/crt.py index 6e8a3f5b8..02182d40c 100644 --- a/packages/smithy-http/src/smithy_http/aio/crt.py +++ b/packages/smithy-http/src/smithy_http/aio/crt.py @@ -1,12 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 # pyright: reportMissingTypeStubs=false,reportUnknownMemberType=false +from asyncio import gather from collections.abc import AsyncGenerator, AsyncIterable -from copy import deepcopy from dataclasses import dataclass from inspect import iscoroutinefunction from io import BytesIO -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Self from awscrt.exceptions import AwsCrtError @@ -195,6 +195,18 @@ async def send( raise _CRTTimeoutError(f"CRT {e.name}: {e.message}") from e raise + async def close(self) -> None: + """Close all pooled HTTP connections.""" + connections = tuple(self._connections.values()) + self._connections.clear() + await gather(*(connection.close() for connection in connections)) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.close() + async def _await_response( self, stream: "AIOHttpClientStreamUnified" ) -> AWSCRTHTTPResponse: @@ -364,7 +376,4 @@ async def _create_body_generator( yield chunk def __deepcopy__(self, memo: Any) -> "AWSCRTHTTPClient": - return AWSCRTHTTPClient( - eventloop=self._eventloop, - client_config=deepcopy(self._config), - ) + return self diff --git a/packages/smithy-http/tests/unit/aio/test_aiohttp.py b/packages/smithy-http/tests/unit/aio/test_aiohttp.py index 4f4548faf..d55846c7f 100644 --- a/packages/smithy-http/tests/unit/aio/test_aiohttp.py +++ b/packages/smithy-http/tests/unit/aio/test_aiohttp.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # pyright: reportPrivateUsage=false from collections.abc import AsyncIterator +from copy import deepcopy from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -17,11 +18,36 @@ def _create_client() -> tuple[AIOHTTPClient, MagicMock]: response.read = AsyncMock(return_value=b"") session = MagicMock() + session.close = AsyncMock() session.request.return_value.__aenter__ = AsyncMock(return_value=response) session.request.return_value.__aexit__ = AsyncMock(return_value=None) return AIOHTTPClient(_session=cast(Any, session)), session +def test_deepcopy_returns_same_client() -> None: + client, _ = _create_client() + + assert deepcopy(client) is client + + +async def test_close_closes_session() -> None: + client, session = _create_client() + + await client.close() + await client.close() + + assert session.close.await_count == 2 + + +async def test_context_manager_closes_session() -> None: + client, session = _create_client() + + async with client as entered: + assert entered is client + + session.close.assert_awaited_once() + + async def test_send_omits_empty_async_reader_body() -> None: client, session = _create_client() request = HTTPRequest( diff --git a/packages/smithy-http/tests/unit/aio/test_crt.py b/packages/smithy-http/tests/unit/aio/test_crt.py index 9718b3a2a..89f6845a0 100644 --- a/packages/smithy-http/tests/unit/aio/test_crt.py +++ b/packages/smithy-http/tests/unit/aio/test_crt.py @@ -22,9 +22,25 @@ def test_deepcopy_client() -> None: - """Test that AWSCRTHTTPClient can be deep copied.""" + """Test that config copies share the stateful HTTP client.""" client = AWSCRTHTTPClient() - deepcopy(client) + assert deepcopy(client) is client + + +async def test_close_closes_and_clears_pooled_connections() -> None: + client = AWSCRTHTTPClient() + connections = [AsyncMock(), AsyncMock()] + client._connections = { + ("https", "one.example.com", None): connections[0], + ("https", "two.example.com", None): connections[1], + } + + await client.close() + await client.close() + + assert client._connections == {} + for connection in connections: + connection.close.assert_awaited_once() def test_client_marshal_request() -> None: