From abbf78fdb9b83885b3e843da426920953c629bee Mon Sep 17 00:00:00 2001
From: Harry Barber <106155934+hlbarber@users.noreply.github.com>
Date: Wed, 22 Mar 2023 15:35:35 +0000
Subject: [PATCH] Make service builder handler setters more flexible (#2442)

* Add generic parameter to the handler setter

* Constrain Pl::Service in Upgrade rather than S

---------

Co-authored-by: Matteo Bigoi <1781140+crisidev@users.noreply.github.com>
---
 .../smithy/generators/ServerServiceGeneratorV2.kt      | 10 +++++-----
 .../aws-smithy-http-server/src/operation/upgrade.rs    |  3 ++-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerServiceGeneratorV2.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerServiceGeneratorV2.kt
index 3b7d09c3c..f5f2cc49d 100644
--- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerServiceGeneratorV2.kt
+++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerServiceGeneratorV2.kt
@@ -129,14 +129,14 @@ class ServerServiceGeneratorV2(
                 /// ## let app: $serviceName<#{SmithyHttpServer}::routing::Route<#{SmithyHttp}::body::SdkBody>> = app;
                 /// ```
                 ///
-                pub fn $fieldName<HandlerType, Extensions>(self, handler: HandlerType) -> Self
+                pub fn $fieldName<HandlerType, HandlerExtractors, ServiceExtractors>(self, handler: HandlerType) -> Self
                 where
-                    HandlerType: #{SmithyHttpServer}::operation::Handler<crate::operation_shape::$structName, Extensions>,
+                    HandlerType: #{SmithyHttpServer}::operation::Handler<crate::operation_shape::$structName, HandlerExtractors>,
                     #{SmithyHttpServer}::operation::Operation<#{SmithyHttpServer}::operation::IntoService<crate::operation_shape::$structName, HandlerType>>:
                         #{SmithyHttpServer}::operation::Upgradable<
                             #{Protocol},
                             crate::operation_shape::$structName,
-                            Extensions,
+                            ServiceExtractors,
                             $builderBodyGenericTypeName,
                             $builderPluginGenericTypeName,
                         >
@@ -151,12 +151,12 @@ class ServerServiceGeneratorV2(
                 /// [`$structName`](crate::operation_shape::$structName) using either
                 /// [`OperationShape::from_handler`](#{SmithyHttpServer}::operation::OperationShapeExt::from_handler) or
                 /// [`OperationShape::from_service`](#{SmithyHttpServer}::operation::OperationShapeExt::from_service).
-                pub fn ${fieldName}_operation<Operation, Extensions>(mut self, operation: Operation) -> Self
+                pub fn ${fieldName}_operation<Operation, Extractors>(mut self, operation: Operation) -> Self
                 where
                     Operation: #{SmithyHttpServer}::operation::Upgradable<
                         #{Protocol},
                         crate::operation_shape::$structName,
-                        Extensions,
+                        Extractors,
                         $builderBodyGenericTypeName,
                         $builderPluginGenericTypeName,
                     >
diff --git a/rust-runtime/aws-smithy-http-server/src/operation/upgrade.rs b/rust-runtime/aws-smithy-http-server/src/operation/upgrade.rs
index 61e550ce8..42c4a5b26 100644
--- a/rust-runtime/aws-smithy-http-server/src/operation/upgrade.rs
+++ b/rust-runtime/aws-smithy-http-server/src/operation/upgrade.rs
@@ -239,7 +239,8 @@ where
     Exts: FromParts<P>,
 
     // The signature of the inner service is correct
-    S: Service<(Op::Input, Exts), Response = Op::Output, Error = OperationError<Op::Error, PollError>> + Clone,
+    Pl::Service:
+        Service<(Op::Input, Exts), Response = Op::Output, Error = OperationError<Op::Error, PollError>> + Clone,
 
     // The plugin takes this operation as input
     Pl: Plugin<P, Op, S, L>,
-- 
GitLab