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 4ab7818b7..b1545202b 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 @@ -5,6 +5,7 @@ package software.amazon.smithy.python.codegen.test; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.nio.file.Files; @@ -46,5 +47,7 @@ public void testCodegen(@TempDir Path tempDir) throws IOException { var client = Files.readString(tempDir.resolve("src/weather/client.py")); assertFalse(client.contains("retry_mode=")); assertFalse(client.contains("max_attempts=")); + assertTrue(client.contains("async def close(self) -> None:")); + assertTrue(client.contains("if self._closed:")); } } 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 7b22daaf2..0c10578b5 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,7 @@ def __init__( self._plugins = plugins self._derive_lock = asyncio.Lock() self._setup_done = False + self._closed = False self._retry_strategy_resolver = $4T() self._client_plugins: list[$2T] = [ ${5C|} @@ -134,6 +135,36 @@ async def _ensure_setup(self) -> None: w.popState(); })); + writer.addStdlibImport("typing", "Any"); + writer.addStdlibImport("typing", "Self"); + writer.write(""" + + async def close(self) -> None: + \"\"\"Close this client and any resources held by its transport.\"\"\" + if self._closed: + return + async with self._derive_lock: + if self._closed: + return + self._closed = True + if self._setup_done and self._config is not None: + await $1T(self._config.transport) + + async def __aenter__(self) -> Self: + if self._closed: + raise RuntimeError("Cannot enter a client that has been closed.") + return self + + async def __aexit__( + self, + exc_type: Any, + exc_value: Any, + traceback: Any, + ) -> None: + await self.close() + """, + RuntimeTypes.ASYNC_CLOSE); + var topDownIndex = TopDownIndex.of(model); var eventStreamIndex = EventStreamIndex.of(model); for (OperationShape operation : topDownIndex.getContainedOperations(service)) { @@ -280,6 +311,11 @@ private void writeSharedOperationInit( writer.write( """ + if self._closed: + raise RuntimeError( + "Cannot invoke an operation on a client that has been closed." + ) + operation_plugins: list[Plugin] = [ $1C ] 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-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 2f3f1c9a0..b2495c118 100644 --- a/packages/smithy-http/src/smithy_http/aio/aiohttp.py +++ b/packages/smithy-http/src/smithy_http/aio/aiohttp.py @@ -2,7 +2,7 @@ # 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 @@ -27,6 +27,7 @@ from smithy_core.interfaces import URI from .. import Field, Fields +from ..exceptions import SmithyHTTPError from ..interfaces import ( HTTPClientConfiguration, HTTPRequestConfiguration, @@ -69,6 +70,7 @@ def __init__( """ _assert_aiohttp() self._config = client_config or AIOHTTPClientConfig() + self._closed = False # Disable transparent response decompression and advertise # 'identity' to request uncompressed responses. # TODO: add a functional test once the test client framework exists @@ -88,6 +90,11 @@ async def send( :param request: The request including destination URI, fields, payload. :param request_config: Configuration specific to this request. """ + if self._closed: + raise SmithyHTTPError( + "Cannot send a request after the HTTP client has been closed." + ) + request_config = request_config or HTTPRequestConfiguration() headers_list = list( @@ -116,6 +123,21 @@ 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.""" + if self._closed: + return + self._closed = True + await self._session.close() + + async def __aenter__(self) -> Self: + if self._closed: + raise SmithyHTTPError("Cannot enter an HTTP client that has been closed.") + 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): diff --git a/packages/smithy-http/src/smithy_http/aio/crt.py b/packages/smithy-http/src/smithy_http/aio/crt.py index ce29313a3..82815b2bd 100644 --- a/packages/smithy-http/src/smithy_http/aio/crt.py +++ b/packages/smithy-http/src/smithy_http/aio/crt.py @@ -1,12 +1,13 @@ # 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 @@ -164,6 +165,7 @@ def __init__( self._tls_ctx = crt_io.ClientTlsContext(crt_io.TlsContextOptions()) self._socket_options = crt_io.SocketOptions() self._connections: ConnectionPoolDict = {} + self._closed = False async def send( self, @@ -176,6 +178,11 @@ async def send( :param request: The request including destination URI, fields, payload. :param request_config: Configuration specific to this request. """ + if self._closed: + raise SmithyHTTPError( + "Cannot send a request after the HTTP client has been closed." + ) + try: crt_request = self._marshal_request(request) connection = await self._get_connection(request.destination) @@ -199,6 +206,26 @@ async def send( raise _CRTTimeoutError(f"CRT {e.name}: {e.message}") from e raise + async def close(self) -> None: + """Close all pooled HTTP connections.""" + if self._closed: + return + self._closed = True + connections = tuple(self._connections.values()) + self._connections.clear() + await gather( + *(connection.close() for connection in connections), + return_exceptions=True, + ) + + async def __aenter__(self) -> Self: + if self._closed: + raise SmithyHTTPError("Cannot enter an HTTP client that has been closed.") + 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: diff --git a/packages/smithy-http/tests/unit/aio/test_aiohttp.py b/packages/smithy-http/tests/unit/aio/test_aiohttp.py index e54233825..c7a9beb25 100644 --- a/packages/smithy-http/tests/unit/aio/test_aiohttp.py +++ b/packages/smithy-http/tests/unit/aio/test_aiohttp.py @@ -5,11 +5,13 @@ from typing import Any, cast from unittest.mock import AsyncMock, MagicMock +import pytest from smithy_core import URI from smithy_core.aio.types import AsyncBytesReader from smithy_http import Field, Fields from smithy_http.aio import HTTPRequest from smithy_http.aio.aiohttp import AIOHTTPClient +from smithy_http.exceptions import SmithyHTTPError def _create_client() -> tuple[AIOHTTPClient, MagicMock]: @@ -17,11 +19,42 @@ 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 +async def test_close_closes_session() -> None: + client, session = _create_client() + + await client.close() + await client.close() + + session.close.assert_awaited_once() + + +async def test_send_after_close_raises() -> None: + client, _ = _create_client() + await client.close() + + with pytest.raises(SmithyHTTPError, match="has been closed"): + await client.send(MagicMock()) + + +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() + + with pytest.raises(SmithyHTTPError, match="has been closed"): + async with client: + pass + + 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 089529b3a..65750b9d3 100644 --- a/packages/smithy-http/tests/unit/aio/test_crt.py +++ b/packages/smithy-http/tests/unit/aio/test_crt.py @@ -27,10 +27,45 @@ def test_deepcopy_client() -> None: deepcopy(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_supports_duplex_streaming() -> None: assert AWSCRTHTTPClient.SUPPORTS_DUPLEX_STREAMING is True +async def test_send_after_close_raises() -> None: + client = AWSCRTHTTPClient() + await client.close() + + with pytest.raises(SmithyHTTPError, match="has been closed"): + await client.send(Mock()) + + +async def test_context_manager_cannot_be_reentered() -> None: + client = AWSCRTHTTPClient() + + async with client as entered: + assert entered is client + + with pytest.raises(SmithyHTTPError, match="has been closed"): + async with client: + pass + + def test_client_marshal_request() -> None: """Test that HTTPRequest is correctly marshaled to CRT HttpRequest.""" client = AWSCRTHTTPClient()