Unify Pokemon model for Rust and Python servers (#2700)

## Motivation and Context
Now we have the feature parity between Rust and Python servers (at least
for the Pokémon service's needs) we can use the same model in both.
Closes https://github.com/awslabs/smithy-rs/issues/1508

## Testing
```bash
$ cd smithy-rs/examples
$ make test # test Rust servers
$ cd python
$ make test # test Python servers
```

## Checklist
<!--- If a checkbox below is not applicable, then please DELETE it
rather than leaving it unchecked -->
- [ ] I have updated `CHANGELOG.next.toml` if I made changes to the
smithy-rs codegen or runtime crates
- [ ] I have updated `CHANGELOG.next.toml` if I made changes to the AWS
SDK, generated SDK code, or SDK runtime crates

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

codegen-client-test/build.gradle.kts

@@ -93,8 +93,8 @@ val allCodegenTests = "../codegen-core/common-test-models".let { commonModels ->
             imports = listOf("$commonModels/naming-obstacle-course-structs.smithy"),
         ),
         CodegenTest("aws.protocoltests.json#TestService", "endpoint-rules"),
-        CodegenTest("com.aws.example.rust#PokemonService", "pokemon-service-client", imports = listOf("$commonModels/pokemon.smithy", "$commonModels/pokemon-common.smithy")),
-        CodegenTest("com.aws.example.rust#PokemonService", "pokemon-service-awsjson-client", imports = listOf("$commonModels/pokemon-awsjson.smithy", "$commonModels/pokemon-common.smithy")),
+        CodegenTest("com.aws.example#PokemonService", "pokemon-service-client", imports = listOf("$commonModels/pokemon.smithy", "$commonModels/pokemon-common.smithy")),
+        CodegenTest("com.aws.example#PokemonService", "pokemon-service-awsjson-client", imports = listOf("$commonModels/pokemon-awsjson.smithy", "$commonModels/pokemon-common.smithy")),
     )
 }
 

codegen-core/common-test-models/pokemon-awsjson.smithy

@@ -4,7 +4,7 @@ $version: "1.0"
 // This is a temporary model to test AwsJson 1.0 with @streaming.
 // This model will be removed when protocol tests support @streaming.
 
-namespace com.aws.example.rust
+namespace com.aws.example
 
 use aws.protocols#awsJson1_0
 use smithy.framework#ValidationException

codegen-core/common-test-models/pokemon.smithy

@@ -1,6 +1,6 @@
 $version: "1.0"
 
-namespace com.aws.example.rust
+namespace com.aws.example
 
 use aws.protocols#restJson1
 use smithy.framework#ValidationException
@@ -20,7 +20,8 @@ service PokemonService {
         GetServerStatistics,
         DoNothing,
         CapturePokemon,
-        CheckHealth
+        CheckHealth,
+        StreamPokemonRadio
     ],
 }
 
@@ -146,3 +147,19 @@ structure MasterBallUnsuccessful {
 
 @error("client")
 structure ThrottlingError {}
+
+/// Fetch a radio song from the database and stream it back as a playable audio.
+@readonly
+@http(uri: "/radio", method: "GET")
+operation StreamPokemonRadio {
+    output: StreamPokemonRadioOutput
+}
+
+@output
+structure StreamPokemonRadioOutput {
+    @httpPayload
+    data: StreamingBlob
+}
+
+@streaming
+blob StreamingBlob

codegen-server-test/build.gradle.kts

@@ -84,12 +84,12 @@ val allCodegenTests = "../codegen-core/common-test-models".let { commonModels ->
         CodegenTest("com.amazonaws.ebs#Ebs", "ebs", imports = listOf("$commonModels/ebs.json")),
         CodegenTest("com.amazonaws.s3#AmazonS3", "s3"),
         CodegenTest(
-            "com.aws.example.rust#PokemonService",
+            "com.aws.example#PokemonService",
             "pokemon-service-server-sdk",
             imports = listOf("$commonModels/pokemon.smithy", "$commonModels/pokemon-common.smithy"),
         ),
         CodegenTest(
-            "com.aws.example.rust#PokemonService",
+            "com.aws.example#PokemonService",
             "pokemon-service-awsjson-server-sdk",
             imports = listOf("$commonModels/pokemon-awsjson.smithy", "$commonModels/pokemon-common.smithy"),
         ),

codegen-server-test/python/build.gradle.kts

@@ -41,7 +41,11 @@ dependencies {
 val allCodegenTests = "../../codegen-core/common-test-models".let { commonModels ->
     listOf(
         CodegenTest("com.amazonaws.simple#SimpleService", "simple", imports = listOf("$commonModels/simple.smithy")),
-        CodegenTest("com.aws.example.python#PokemonService", "pokemon-service-server-sdk"),
+        CodegenTest(
+            "com.aws.example#PokemonService",
+            "pokemon-service-server-sdk",
+            imports = listOf("$commonModels/pokemon.smithy", "$commonModels/pokemon-common.smithy"),
+        ),
         CodegenTest(
             "com.amazonaws.ebs#Ebs", "ebs",
             imports = listOf("$commonModels/ebs.json"),

codegen-server-test/python/model/pokemon-common.smithy

@@ -1 +0,0 @@
-../../../codegen-core/common-test-models/pokemon-common.smithy

codegen-server-test/python/model/pokemon.smithy

@@ -1,126 +0,0 @@
-/// TODO(https://github.com/awslabs/smithy-rs/issues/1508)
-/// Reconcile this model with the main one living inside codegen-server-test/model/pokemon.smithy
-/// once the Python implementation supports Streaming and Union shapes.
-$version: "1.0"
-
-namespace com.aws.example.python
-
-use aws.protocols#restJson1
-use com.aws.example#CheckHealth
-use com.aws.example#DoNothing
-use com.aws.example#GetServerStatistics
-use com.aws.example#PokemonSpecies
-use com.aws.example#Storage
-use smithy.framework#ValidationException
-
-/// The Pokémon Service allows you to retrieve information about Pokémon species.
-@title("Pokémon Service")
-@restJson1
-service PokemonService {
-    version: "2021-12-01"
-    resources: [PokemonSpecies]
-    operations: [
-        GetServerStatistics
-        DoNothing
-        CapturePokemon
-        CheckHealth
-        StreamPokemonRadio
-    ]
-}
-
-/// Capture Pokémons via event streams.
-@http(uri: "/capture-pokemon-event/{region}", method: "POST")
-operation CapturePokemon {
-    input: CapturePokemonEventsInput
-    output: CapturePokemonEventsOutput
-    errors: [
-        UnsupportedRegionError
-        ThrottlingError
-        ValidationException
-    ]
-}
-
-@input
-structure CapturePokemonEventsInput {
-    @httpPayload
-    events: AttemptCapturingPokemonEvent
-    @httpLabel
-    @required
-    region: String
-}
-
-@output
-structure CapturePokemonEventsOutput {
-    @httpPayload
-    events: CapturePokemonEvents
-}
-
-@streaming
-union AttemptCapturingPokemonEvent {
-    event: CapturingEvent
-    masterball_unsuccessful: MasterBallUnsuccessful
-}
-
-structure CapturingEvent {
-    @eventPayload
-    payload: CapturingPayload
-}
-
-structure CapturingPayload {
-    name: String
-    pokeball: String
-}
-
-@streaming
-union CapturePokemonEvents {
-    event: CaptureEvent
-    invalid_pokeball: InvalidPokeballError
-    throttlingError: ThrottlingError
-}
-
-structure CaptureEvent {
-    @eventHeader
-    name: String
-    @eventHeader
-    captured: Boolean
-    @eventHeader
-    shiny: Boolean
-    @eventPayload
-    pokedex_update: Blob
-}
-
-@error("server")
-structure UnsupportedRegionError {
-    @required
-    region: String
-}
-
-@error("client")
-structure InvalidPokeballError {
-    @required
-    pokeball: String
-}
-
-@error("server")
-structure MasterBallUnsuccessful {
-    message: String
-}
-
-@error("client")
-structure ThrottlingError {}
-
-/// Fetch the radio song from the database and stream it back as a playable audio.
-@readonly
-@http(uri: "/radio", method: "GET")
-operation StreamPokemonRadio {
-    output: StreamPokemonRadioOutput
-}
-
-@output
-structure StreamPokemonRadioOutput {
-    @httpPayload
-    data: StreamingBlob
-}
-
-@streaming
-blob StreamingBlob

examples/pokemon-service-common/Cargo.toml

@@ -13,12 +13,11 @@ rand = "0.8"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3.16", features = ["env-filter", "json"] }
 tokio = { version = "1", default-features = false, features = ["time"] }
+tower = "0.4"
 
 # Local paths
+aws-smithy-client = { path = "../../rust-runtime/aws-smithy-client" }
 aws-smithy-http = { path = "../../rust-runtime/aws-smithy-http" }
 aws-smithy-http-server = { path = "../../rust-runtime/aws-smithy-http-server" }
 pokemon-service-client = { path = "../pokemon-service-client" }
 pokemon-service-server-sdk = { path = "../pokemon-service-server-sdk" }
-
-[dev-dependencies]
-tower = "0.4"

examples/pokemon-service-common/src/lib.rs

@@ -15,7 +15,8 @@ use std::{
 };
 
 use async_stream::stream;
-use aws_smithy_http::operation::Request;
+use aws_smithy_client::{conns, hyper_ext::Adapter};
+use aws_smithy_http::{body::SdkBody, byte_stream::ByteStream, operation::Request};
 use aws_smithy_http_server::Extension;
 use http::{
     uri::{Authority, Scheme},
@@ -24,7 +25,8 @@ use http::{
 use pokemon_service_server_sdk::{
     error, input, model, model::CapturingPayload, output, types::Blob,
 };
-use rand::Rng;
+use rand::{seq::SliceRandom, Rng};
+use tower::Service;
 use tracing_subscriber::{prelude::*, EnvFilter};
 
 const PIKACHU_ENGLISH_FLAVOR_TEXT: &str =
@@ -327,6 +329,37 @@ pub async fn check_health(_input: input::CheckHealthInput) -> output::CheckHealt
     output::CheckHealthOutput {}
 }
 
+const RADIO_STREAMS: [&str; 2] = [
+    "https://ia800107.us.archive.org/33/items/299SoundEffectCollection/102%20Palette%20Town%20Theme.mp3",
+    "https://ia600408.us.archive.org/29/items/PocketMonstersGreenBetaLavenderTownMusicwwwFlvtoCom/Pocket%20Monsters%20Green%20Beta-%20Lavender%20Town%20Music-%5Bwww_flvto_com%5D.mp3",
+];
+
+/// Streams a random Pokémon song.
+pub async fn stream_pokemon_radio(
+    _input: input::StreamPokemonRadioInput,
+) -> output::StreamPokemonRadioOutput {
+    let radio_stream_url = RADIO_STREAMS
+        .choose(&mut rand::thread_rng())
+        .expect("`RADIO_STREAMS` is empty")
+        .parse::<Uri>()
+        .expect("Invalid url in `RADIO_STREAMS`");
+
+    let mut connector = Adapter::builder().build(conns::https());
+    let result = connector
+        .call(
+            http::Request::builder()
+                .uri(radio_stream_url)
+                .body(SdkBody::empty())
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    output::StreamPokemonRadioOutput {
+        data: ByteStream::new(result.into_body()),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

examples/pokemon-service-lambda/src/main.rs

@@ -9,7 +9,7 @@ use aws_smithy_http_server::{routing::LambdaHandler, AddExtensionLayer};
 
 use pokemon_service_common::{
     capture_pokemon, check_health, do_nothing, get_pokemon_species, get_server_statistics,
-    setup_tracing, State,
+    setup_tracing, stream_pokemon_radio, State,
 };
 use pokemon_service_lambda::get_storage_lambda;
 use pokemon_service_server_sdk::PokemonService;
@@ -28,6 +28,7 @@ pub async fn main() {
         .capture_pokemon(capture_pokemon)
         .do_nothing(do_nothing)
         .check_health(check_health)
+        .stream_pokemon_radio(stream_pokemon_radio)
         .build()
         .expect("failed to build an instance of PokemonService")
         // Set up shared state and middlewares.

examples/pokemon-service-tls/src/main.rs

@@ -34,7 +34,7 @@ use tokio_rustls::{
 
 use pokemon_service_common::{
     capture_pokemon, check_health, do_nothing, get_pokemon_species, get_server_statistics,
-    get_storage, setup_tracing, State,
+    get_storage, setup_tracing, stream_pokemon_radio, State,
 };
 use pokemon_service_server_sdk::PokemonService;
 use pokemon_service_tls::{DEFAULT_ADDRESS, DEFAULT_PORT, DEFAULT_TEST_CERT, DEFAULT_TEST_KEY};
@@ -71,6 +71,7 @@ pub async fn main() {
         .capture_pokemon(capture_pokemon)
         .do_nothing(do_nothing)
         .check_health(check_health)
+        .stream_pokemon_radio(stream_pokemon_radio)
         .build()
         .expect("failed to build an instance of PokemonService")
         // Set up shared state and middlewares.

examples/pokemon-service/src/main.rs

@@ -23,7 +23,8 @@ use pokemon_service::{
     do_nothing_but_log_request_ids, get_storage_with_local_approved, DEFAULT_ADDRESS, DEFAULT_PORT,
 };
 use pokemon_service_common::{
-    capture_pokemon, check_health, get_pokemon_species, get_server_statistics, setup_tracing, State,
+    capture_pokemon, check_health, get_pokemon_species, get_server_statistics, setup_tracing,
+    stream_pokemon_radio, State,
 };
 use pokemon_service_server_sdk::PokemonService;
 
@@ -67,6 +68,7 @@ pub async fn main() {
         .capture_pokemon(capture_pokemon)
         .do_nothing(do_nothing_but_log_request_ids)
         .check_health(check_health)
+        .stream_pokemon_radio(stream_pokemon_radio)
         .build()
         .expect("failed to build an instance of PokemonService");
 

examples/python/Makefile

@@ -45,7 +45,10 @@ release: codegen
 	$(MAKE) generate-stubs
 	$(MAKE) build-wheel-release
 
-run: build
+run: build install-wheel
+	python3 $(CUR_DIR)/pokemon_service.py
+
+run-release: release install-wheel
 	python3 $(CUR_DIR)/pokemon_service.py
 
 py-check: install-wheel

examples/python/pokemon_service.py

@@ -17,8 +17,10 @@ from pokemon_service_server_sdk.error import (
     MasterBallUnsuccessful,
     ResourceNotFoundException,
     UnsupportedRegionError,
+    StorageAccessNotAuthorized,
 )
 from pokemon_service_server_sdk.input import (
+    GetStorageInput,
     CapturePokemonInput,
     CheckHealthInput,
     DoNothingInput,
@@ -28,8 +30,14 @@ from pokemon_service_server_sdk.input import (
 )
 from pokemon_service_server_sdk.logging import TracingHandler
 from pokemon_service_server_sdk.middleware import MiddlewareException, Request, Response
-from pokemon_service_server_sdk.model import CaptureEvent, CapturePokemonEvents, FlavorText, Language
+from pokemon_service_server_sdk.model import (
+    CaptureEvent,
+    CapturePokemonEvents,
+    FlavorText,
+    Language,
+)
 from pokemon_service_server_sdk.output import (
+    GetStorageOutput,
     CapturePokemonOutput,
     CheckHealthOutput,
     DoNothingOutput,
@@ -164,7 +172,11 @@ async def check_content_type_header(request: Request, next: Next) -> Response:
     if content_type == "application/json":
         logging.debug("found valid `application/json` content type")
     else:
-        logging.warning("invalid content type %s, dumping headers: %s", content_type, request.headers.items())
+        logging.warning(
+            "invalid content type %s, dumping headers: %s",
+            content_type,
+            request.headers.items(),
+        )
     return await next(request)
 
 
@@ -198,9 +210,24 @@ def do_nothing(_: DoNothingInput) -> DoNothingOutput:
     return DoNothingOutput()
 
 
+# Retrieves the user's storage.
+@app.get_storage
+def get_storage(input: GetStorageInput) -> GetStorageOutput:
+    logging.debug("attempting to authenticate storage user")
+
+    # We currently only support Ash and he has nothing stored
+    if input.user != "ash" or input.passcode != "pikachu123":
+        logging.debug("authentication failed")
+        raise StorageAccessNotAuthorized()
+
+    return GetStorageOutput([])
+
+
 # Get the translation of a Pokémon specie or an error.
 @app.get_pokemon_species
-def get_pokemon_species(input: GetPokemonSpeciesInput, context: Context) -> GetPokemonSpeciesOutput:
+def get_pokemon_species(
+    input: GetPokemonSpeciesInput, context: Context
+) -> GetPokemonSpeciesOutput:
     if context.lambda_ctx is not None:
         logging.debug(
             "lambda Context: %s",
@@ -218,7 +245,9 @@ def get_pokemon_species(input: GetPokemonSpeciesInput, context: Context) -> GetP
     if flavor_text_entries:
         logging.debug("total requests executed: %s", context.get_calls_count())
         logging.info("found description for Pokémon %s", input.name)
-        return GetPokemonSpeciesOutput(name=input.name, flavor_text_entries=flavor_text_entries)
+        return GetPokemonSpeciesOutput(
+            name=input.name, flavor_text_entries=flavor_text_entries
+        )
     else:
         logging.warning("description for Pokémon %s not in the database", input.name)
         raise ResourceNotFoundException("Requested Pokémon not available")
@@ -226,7 +255,9 @@ def get_pokemon_species(input: GetPokemonSpeciesInput, context: Context) -> GetP
 
 # Get the number of requests served by this server.
 @app.get_server_statistics
-def get_server_statistics(_: GetServerStatisticsInput, context: Context) -> GetServerStatisticsOutput:
+def get_server_statistics(
+    _: GetServerStatisticsInput, context: Context
+) -> GetServerStatisticsOutput:
     calls_count = context.get_calls_count()
     logging.debug("the service handled %d requests", calls_count)
     return GetServerStatisticsOutput(calls_count=calls_count)
@@ -393,7 +424,9 @@ def capture_pokemon(input: CapturePokemonInput) -> CapturePokemonOutput:
 
 # Stream a random Pokémon song.
 @app.stream_pokemon_radio
-async def stream_pokemon_radio(_: StreamPokemonRadioInput, context: Context) -> StreamPokemonRadioOutput:
+async def stream_pokemon_radio(
+    _: StreamPokemonRadioInput, context: Context
+) -> StreamPokemonRadioOutput:
     import aiohttp
 
     radio_url = context.get_random_radio_stream()