Unverified Commit 7c9c283a authored by ysaito1001's avatar ysaito1001 Committed by GitHub
Browse files

Add `Metadata` to a config layer through operation runtime plugin (#2830) 

## Motivation and Context
Adds `Metadata` to a config layer through `OperationRuntimePlugin`.

## Description
We have had a customer's use case where a service name and a operation
name are obtained from
[Metadata](https://github.com/awslabs/smithy-rs/blob/ddba46086a9754c01f3b11d7521c49d4489de84b/rust-runtime/aws-smithy-http/src/operation.rs#L17-L22

).
The end goal was to use their names as part of metrics collection.
Previously, it was done using `map_operation` on a
`CustomizableOperation`, e.g.
```
client
    .some_operation()
    .customize()
    .await?
    .map_operation(|operation| {
        operation.try_clone().map(|operation| {
            let (_, parts) = operation.into_request_response();
            parts.metadata.map(|metadata| {
                let service_name = metadata.service().to_string().to_uppercase();
                let operation_name = metadata.name().to_string();
                /*
                 * do something with `service_name` and `operation_name`
                 */
            })
        });

        Ok(operation)
    })?
    .send()
    .await;
```
The orchestrator no longer supports `map_operation` on
`CustomizableOperation`. We therefore add `Metadata` to a config layer
through `OperationRuntimePlugin`. See the added integration test for how
`Metadata` is retrieved from within an interceptor.

## Testing
Added an integration-test to verify `Metadata` is properly set.

----

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

---------

Co-authored-by: default avatarysaito1001 <awsaito@amazon.com>
Co-authored-by: default avatarZelda Hessler <zhessler@amazon.com>
parent 8abc9463

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/MetadataCustomization.kt

0 → 100644
+48 −0
Original line number Diff line number Diff line 
/*

 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

 * SPDX-License-Identifier: Apache-2.0

 */



package software.amazon.smithy.rust.codegen.client.smithy.customizations



import software.amazon.smithy.model.shapes.OperationShape

import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext

import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCustomization

import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection

import software.amazon.smithy.rust.codegen.core.rustlang.Writable

import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate

import software.amazon.smithy.rust.codegen.core.rustlang.writable

import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType

import software.amazon.smithy.rust.codegen.core.util.dq

import software.amazon.smithy.rust.codegen.core.util.sdkId



class MetadataCustomization(

    private val codegenContext: ClientCodegenContext,

    operation: OperationShape,

) : OperationCustomization() {

    private val operationName = codegenContext.symbolProvider.toSymbol(operation).name

    private val runtimeConfig = codegenContext.runtimeConfig

    private val codegenScope by lazy {

        arrayOf(

            "Metadata" to RuntimeType.operationModule(runtimeConfig).resolve("Metadata"),

        )

    }



    override fun section(section: OperationSection): Writable = writable {

        when (section) {

            is OperationSection.AdditionalRuntimePluginConfig -> {

                rustTemplate(

                    """

                    ${section.newLayerName}.store_put(#{Metadata}::new(

                        ${operationName.dq()},

                        ${codegenContext.serviceShape.sdkId().dq()},

                    ));

                    """,

                    *codegenScope,

                )

            }



            else -> {}

        }

    }

}

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/RequiredCustomizations.kt

+5 −1
Original line number Diff line number Diff line
@@ -15,6 +15,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customizations.HttpVers 
import software.amazon.smithy.rust.codegen.client.smithy.customizations.IdempotencyTokenGenerator

import software.amazon.smithy.rust.codegen.client.smithy.customizations.IdentityConfigCustomization

import software.amazon.smithy.rust.codegen.client.smithy.customizations.InterceptorConfigCustomization

import software.amazon.smithy.rust.codegen.client.smithy.customizations.MetadataCustomization

import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyConfigCustomization

import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyReExportCustomization

import software.amazon.smithy.rust.codegen.client.smithy.customizations.ResiliencyServiceRuntimePluginCustomization
@@ -30,6 +31,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.customizations.CrateVersi 
import software.amazon.smithy.rust.codegen.core.smithy.customizations.pubUseSmithyErrorTypes

import software.amazon.smithy.rust.codegen.core.smithy.customizations.pubUseSmithyPrimitives

import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization

import software.amazon.smithy.rust.codegen.core.util.letIf



val TestUtilFeature = Feature("test-util", false, listOf())
@@ -47,7 +49,9 @@ class RequiredCustomizations : ClientCodegenDecorator { 
        operation: OperationShape,

        baseCustomizations: List<OperationCustomization>,

    ): List<OperationCustomization> =

        baseCustomizations +

        baseCustomizations.letIf(codegenContext.smithyRuntimeMode.generateOrchestrator) {

            it + MetadataCustomization(codegenContext, operation)

        } +

            IdempotencyTokenGenerator(codegenContext, operation) +

            EndpointPrefixGenerator(codegenContext, operation) +

            HttpChecksumRequiredGenerator(codegenContext, operation) +

codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/MetadataCustomizationTest.kt

0 → 100644
+113 −0
Original line number Diff line number Diff line 
/*

 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

 * SPDX-License-Identifier: Apache-2.0

 */



package software.amazon.smithy.rust.codegen.client.smithy.customizations



import org.junit.jupiter.api.Test

import software.amazon.smithy.rust.codegen.client.testutil.TestCodegenSettings

import software.amazon.smithy.rust.codegen.client.testutil.clientIntegrationTest

import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency

import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate

import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType

import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope

import software.amazon.smithy.rust.codegen.core.testutil.IntegrationTestParams

import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel

import software.amazon.smithy.rust.codegen.core.testutil.testModule

import software.amazon.smithy.rust.codegen.core.testutil.tokioTest



class MetadataCustomizationTest {

    private val model = """

        namespace com.example

        use aws.protocols#awsJson1_0

        @awsJson1_0

        service HelloService {

            operations: [SayHello],

            version: "1"

        }

        @optionalAuth

        operation SayHello { input: TestInput }

        structure TestInput {

           foo: String,

        }

    """.asSmithyModel()



    @Test

    fun `extract metadata via customizable operation`() {

        clientIntegrationTest(

            model,

            params = IntegrationTestParams(additionalSettings = TestCodegenSettings.orchestratorMode()),

        ) { clientCodegenContext, rustCrate ->

            val runtimeConfig = clientCodegenContext.runtimeConfig

            val codegenScope = arrayOf(

                *preludeScope,

                "BeforeTransmitInterceptorContextMut" to RuntimeType.beforeTransmitInterceptorContextMut(runtimeConfig),

                "BoxError" to RuntimeType.boxError(runtimeConfig),

                "ConfigBag" to RuntimeType.configBag(runtimeConfig),

                "Interceptor" to RuntimeType.interceptor(runtimeConfig),

                "Metadata" to RuntimeType.operationModule(runtimeConfig).resolve("Metadata"),

                "capture_request" to RuntimeType.captureRequest(runtimeConfig),

            )

            rustCrate.testModule {

                addDependency(CargoDependency.Tokio.withFeature("test-util").toDevDependency())

                tokioTest("test_extract_metadata_via_customizable_operation") {

                    rustTemplate(

                        """

                        // Interceptors aren’t supposed to store states, but it is done this way for a testing purpose.

                        ##[derive(Debug)]

                        struct ExtractMetadataInterceptor(

                            ::std::sync::Mutex<#{Option}<::std::sync::mpsc::Sender<(String, String)>>>,

                        );



                        impl #{Interceptor} for ExtractMetadataInterceptor {

                            fn modify_before_signing(

                                &self,

                                _context: &mut #{BeforeTransmitInterceptorContextMut}<'_>,

                                cfg: &mut #{ConfigBag},

                            ) -> #{Result}<(), #{BoxError}> {

                                let metadata = cfg

                                    .load::<#{Metadata}>()

                                    .expect("metadata should exist");

                                let service_name = metadata.service().to_string();

                                let operation_name = metadata.name().to_string();

                                let tx = self.0.lock().unwrap().take().unwrap();

                                tx.send((service_name, operation_name)).unwrap();

                                #{Ok}(())

                            }

                        }



                        let (tx, rx) = ::std::sync::mpsc::channel();



                        let (conn, _captured_request) = #{capture_request}(#{None});

                        let client_config = crate::config::Config::builder()

                            .endpoint_resolver("http://localhost:1234/")

                            .http_connector(conn)

                            .build();

                        let client = crate::client::Client::from_conf(client_config);

                        let _ = client

                            .say_hello()

                            .customize()

                            .await

                            .expect("operation should be customizable")

                            .interceptor(ExtractMetadataInterceptor(::std::sync::Mutex::new(#{Some}(tx))))

                            .send()

                            .await;



                        match rx.recv() {

                            #{Ok}((service_name, operation_name)) => {

                                assert_eq!("HelloService", &service_name);

                                assert_eq!("SayHello", &operation_name);

                            }

                            #{Err}(_) => panic!(

                                "failed to receive service name and operation name from `ExtractMetadataInterceptor`"

                            ),

                        }

                        """,

                        *codegenScope,

                    )

                }

            }

        }

    }

}

rust-runtime/aws-smithy-http/src/operation.rs

+5 −0
Original line number Diff line number Diff line
@@ -9,6 +9,7 @@ 
use crate::body::SdkBody;

use crate::property_bag::{PropertyBag, SharedPropertyBag};

use crate::retry::DefaultResponseRetryClassifier;

use aws_smithy_types::config_bag::{Storable, StoreReplace};

use std::borrow::Cow;

use std::ops::{Deref, DerefMut};
@@ -44,6 +45,10 @@ impl Metadata { 
    }

}



impl Storable for Metadata {

    type Storer = StoreReplace<Self>;

}



/// Non-request parts of an [`Operation`].

///

/// Generics: