Unverified Commit a514793f authored by John DiSanti's avatar John DiSanti Committed by GitHub
Browse files

Port API Gateway customizations to the orchestrator (#2705)

## Motivation and Context
This PR ports the API Gateway `Accept` header customization into the
orchestrator implementation. The original implementation of this
customization in #287 didn't include a test, so this PR also adds an
integration test for it.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
parent 3b95b970
Loading
Loading
Loading
Loading
+1 −0
Original line number Diff line number Diff line
@@ -20,6 +20,7 @@ aws-smithy-checksums = { path = "../../../rust-runtime/aws-smithy-checksums" }
aws-smithy-client = { path = "../../../rust-runtime/aws-smithy-client" }
aws-smithy-http = { path = "../../../rust-runtime/aws-smithy-http" }
aws-smithy-http-tower= { path = "../../../rust-runtime/aws-smithy-http-tower" }
aws-smithy-runtime-api = { path = "../../../rust-runtime/aws-smithy-runtime-api" }
aws-smithy-types = { path = "../../../rust-runtime/aws-smithy-types" }
aws-types = { path = "../aws-types" }
bytes = "1"
+30 −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
 */

#![allow(dead_code)]

use aws_smithy_runtime_api::client::interceptors::context::phase::BeforeTransmit;
use aws_smithy_runtime_api::client::interceptors::{BoxError, Interceptor, InterceptorContext};
use aws_smithy_runtime_api::config_bag::ConfigBag;
use http::header::ACCEPT;
use http::HeaderValue;

/// Interceptor that adds an Accept header to API Gateway requests.
#[derive(Debug, Default)]
pub(crate) struct AcceptHeaderInterceptor;

impl Interceptor for AcceptHeaderInterceptor {
    fn modify_before_signing(
        &self,
        context: &mut InterceptorContext<BeforeTransmit>,
        _cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        context
            .request_mut()
            .headers_mut()
            .insert(ACCEPT, HeaderValue::from_static("application/json"));
        Ok(())
    }
}
+3 −0
Original line number Diff line number Diff line
@@ -19,6 +19,9 @@
    unreachable_pub
)]

/// Interceptors for API Gateway
pub mod apigateway_interceptors;

/// Stub credentials provider for use when no credentials provider is used.
pub mod no_credentials;

+8 −0
Original line number Diff line number Diff line
@@ -14,6 +14,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientProtocolMap
import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointCustomization
import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization
import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
import software.amazon.smithy.rust.codegen.client.smithy.generators.error.ErrorCustomization
import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
@@ -130,4 +131,11 @@ class ServiceSpecificDecorator(
        model.maybeApply(service) {
            delegateTo.transformModel(service, model)
        }

    override fun serviceRuntimePluginCustomizations(
        codegenContext: ClientCodegenContext,
        baseCustomizations: List<ServiceRuntimePluginCustomization>,
    ): List<ServiceRuntimePluginCustomization> = baseCustomizations.maybeApply(codegenContext.serviceShape) {
        delegateTo.serviceRuntimePluginCustomizations(codegenContext, baseCustomizations)
    }
}
+43 −2
Original line number Diff line number Diff line
@@ -8,25 +8,44 @@ package software.amazon.smithy.rustsdk.customize.apigateway
import software.amazon.smithy.model.shapes.OperationShape
import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginCustomization
import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceRuntimePluginSection
import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
import software.amazon.smithy.rust.codegen.core.rustlang.Writable
import software.amazon.smithy.rust.codegen.core.rustlang.rust
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.smithy.customize.OperationCustomization
import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationSection
import software.amazon.smithy.rust.codegen.core.util.letIf
import software.amazon.smithy.rustsdk.InlineAwsDependency

class ApiGatewayDecorator : ClientCodegenDecorator {
    override val name: String = "ApiGateway"
    override val order: Byte = 0

    // TODO(enableNewSmithyRuntime): Delete when cleaning up middleware
    override fun operationCustomizations(
        codegenContext: ClientCodegenContext,
        operation: OperationShape,
        baseCustomizations: List<OperationCustomization>,
    ): List<OperationCustomization> = baseCustomizations + ApiGatewayAddAcceptHeader()
    ): List<OperationCustomization> =
        baseCustomizations.letIf(codegenContext.smithyRuntimeMode.generateMiddleware) {
            it + ApiGatewayAddAcceptHeader()
        }

class ApiGatewayAddAcceptHeader : OperationCustomization() {
    override fun serviceRuntimePluginCustomizations(
        codegenContext: ClientCodegenContext,
        baseCustomizations: List<ServiceRuntimePluginCustomization>,
    ): List<ServiceRuntimePluginCustomization> =
        baseCustomizations.letIf(codegenContext.smithyRuntimeMode.generateOrchestrator) {
            it + ApiGatewayAcceptHeaderInterceptorCustomization(codegenContext)
        }
}

// TODO(enableNewSmithyRuntime): Delete when cleaning up middleware
private class ApiGatewayAddAcceptHeader : OperationCustomization() {
    override fun section(section: OperationSection): Writable = when (section) {
        is OperationSection.FinalizeOperation -> emptySection
        is OperationSection.OperationImplBlock -> emptySection
@@ -39,6 +58,28 @@ class ApiGatewayAddAcceptHeader : OperationCustomization() {
                RuntimeType.Http,
            )
        }

        else -> emptySection
    }
}

private class ApiGatewayAcceptHeaderInterceptorCustomization(private val codegenContext: ClientCodegenContext) :
    ServiceRuntimePluginCustomization() {
    override fun section(section: ServiceRuntimePluginSection): Writable = writable {
        if (section is ServiceRuntimePluginSection.AdditionalConfig) {
            section.registerInterceptor(codegenContext.runtimeConfig, this) {
                rustTemplate(
                    "#{Interceptor}::default()",
                    "Interceptor" to RuntimeType.forInlineDependency(
                        InlineAwsDependency.forRustFile(
                            "apigateway_interceptors",
                            additionalDependency = arrayOf(
                                CargoDependency.smithyRuntimeApi(codegenContext.runtimeConfig),
                            ),
                        ),
                    ).resolve("AcceptHeaderInterceptor"),
                )
            }
        }
    }
}
Loading