Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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__":
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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())
Expand All @@ -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}"));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand Down Expand Up @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<String> CLIENT_RESERVED_METHOD_NAMES = Set.of("close");

private final Model model;
private final ReservedWordSymbolProvider.Escaper escaper;
Expand Down Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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();
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type": "bugfix",
"description": "Preserved HTTP clients across operation configuration copies to avoid duplicating sessions and discarding connection pools."
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type": "feature",
"description": "Added deterministic connection-pool cleanup and async context manager support to HTTP clients."
}
18 changes: 12 additions & 6 deletions packages/smithy-http/src/smithy_http/aio/aiohttp.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
21 changes: 15 additions & 6 deletions packages/smithy-http/src/smithy_http/aio/crt.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
26 changes: 26 additions & 0 deletions packages/smithy-http/tests/unit/aio/test_aiohttp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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(
Expand Down
20 changes: 18 additions & 2 deletions packages/smithy-http/tests/unit/aio/test_crt.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading