Python: Allow injecting Lambda Context (#1985)

* Provide Python wrappers for Lambda related types

* Introduce `PyContext` to wrap raw context object

* Use new `PyContext` in handlers

* Expose `lambda` module to Python

* Use `LambdaContext` in example service

* Start Lambda handler in a different thread

* Print summary of Lambda context in Pokemon service

* Make sure to include Python `builtins` in tests

* Make `lambda_ctx` optional

Co-authored-by: Matteo Bigoi <1781140+crisidev@users.noreply.github.com>

* Only inject types if they are type-hinted as `Optional[T]`

* Export Lambda module as `aws_lambda` instead of `lambda_`

* Comment why we need to run Hyper server in a background thread

* Move `is_optional_of` to `util` module

* Use `HeaderMap::from_iter` to build headers

* Support edge case of `(None, T)` in `util::is_optional_of`

* Make Lambda related types feature gated

* Remove feature gate for Lambda

* Make `xray_trace_id` an `Option`

* Remove `aws-lambda` feature from generated `Cargo.toml`s

* Fix linting issues

* Pin `lambda_runtime` to `0.7.1`

* Remove duplicate dependency in `Cargo.toml`

Co-authored-by: Matteo Bigoi <1781140+crisidev@users.noreply.github.com>

Author: unexge

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/customizations/PythonServerCodegenDecorator.kt

@@ -7,7 +7,6 @@ package software.amazon.smithy.rust.codegen.server.python.smithy.customizations
 
 import software.amazon.smithy.model.neighbor.Walker
 import software.amazon.smithy.rust.codegen.client.smithy.customize.RustCodegenDecorator
-import software.amazon.smithy.rust.codegen.core.rustlang.Feature
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.docs
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
@@ -104,21 +103,6 @@ class PubUsePythonTypesDecorator : RustCodegenDecorator<ServerProtocolGenerator,
         clazz.isAssignableFrom(ServerCodegenContext::class.java)
 }
 
-/**
- * Decorator adding an `aws-lambda` feature to the generated crate.
- */
-class PythonFeatureFlagsDecorator : RustCodegenDecorator<ServerProtocolGenerator, ServerCodegenContext> {
-    override val name: String = "PythonFeatureFlagsDecorator"
-    override val order: Byte = 0
-
-    override fun extras(codegenContext: ServerCodegenContext, rustCrate: RustCrate) {
-        rustCrate.mergeFeature(Feature("aws-lambda", true, listOf("aws-smithy-http-server-python/aws-lambda")))
-    }
-
-    override fun supportsCodegenContext(clazz: Class<out CodegenContext>): Boolean =
-        clazz.isAssignableFrom(ServerCodegenContext::class.java)
-}
-
 val DECORATORS = listOf(
     /**
      * Add the [InternalServerError] error to all operations.
@@ -131,6 +115,4 @@ val DECORATORS = listOf(
     PubUsePythonTypesDecorator(),
     // Render the Python shared library export.
     PythonExportModuleDecorator(),
-    // Add the `aws-lambda` feature flag
-    PythonFeatureFlagsDecorator(),
 )

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/generators/PythonServerModuleGenerator.kt

@@ -49,6 +49,7 @@ class PythonServerModuleGenerator(
                 renderPyLogging()
                 renderPyMiddlewareTypes()
                 renderPyTlsTypes()
+                renderPyLambdaTypes()
                 renderPyApplicationType()
             }
         }
@@ -179,6 +180,22 @@ class PythonServerModuleGenerator(
         )
     }
 
+    private fun RustWriter.renderPyLambdaTypes() {
+        rustTemplate(
+            """
+            let aws_lambda = #{pyo3}::types::PyModule::new(py, "aws_lambda")?;
+            aws_lambda.add_class::<#{SmithyPython}::lambda::PyLambdaContext>()?;
+            pyo3::py_run!(
+                py,
+                aws_lambda,
+                "import sys; sys.modules['$libName.aws_lambda'] = aws_lambda"
+            );
+            m.add_submodule(aws_lambda)?;
+            """,
+            *codegenScope,
+        )
+    }
+
     // Render Python application type.
     private fun RustWriter.renderPyApplicationType() {
         rustTemplate(

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/generators/PythonServerOperationHandlerGenerator.kt

@@ -65,7 +65,7 @@ class PythonServerOperationHandlerGenerator(
                 /// Python handler for operation `$operationName`.
                 pub(crate) async fn $fnName(
                     input: $input,
-                    state: #{SmithyServer}::Extension<#{pyo3}::PyObject>,
+                    state: #{SmithyServer}::Extension<#{SmithyPython}::context::PyContext>,
                     handler: #{SmithyPython}::PyHandler,
                 ) -> std::result::Result<$output, $error> {
                     // Async block used to run the handler and catch any Python error.
@@ -95,7 +95,7 @@ class PythonServerOperationHandlerGenerator(
                     let output = if handler.args == 1 {
                         pyhandler.call1((input,))?
                     } else {
-                        pyhandler.call1((input, state.0))?
+                        pyhandler.call1((input, #{pyo3}::ToPyObject::to_object(&state.0, py)))?
                     };
                     output.extract::<$output>()
                 })
@@ -114,7 +114,7 @@ class PythonServerOperationHandlerGenerator(
                     let coroutine = if handler.args == 1 {
                         pyhandler.call1((input,))?
                     } else {
-                        pyhandler.call1((input, state.0))?
+                        pyhandler.call1((input, #{pyo3}::ToPyObject::to_object(&state.0, py)))?
                     };
                     #{pyo3_asyncio}::tokio::into_future(coroutine)
                 })?;

rust-runtime/aws-smithy-http-server-python/Cargo.toml

@@ -12,23 +12,25 @@ Python server runtime for Smithy Rust Server Framework.
 """
 publish = true
 
-[features]
-aws-lambda = ["aws-smithy-http-server/aws-lambda", "dep:lambda_http"]
-
 [dependencies]
 aws-smithy-http = { path = "../aws-smithy-http" }
-aws-smithy-http-server = { path = "../aws-smithy-http-server" }
+aws-smithy-http-server = { path = "../aws-smithy-http-server", features = ["aws-lambda"] }
 aws-smithy-json = { path = "../aws-smithy-json" }
 aws-smithy-types = { path = "../aws-smithy-types" }
 aws-smithy-xml = { path = "../aws-smithy-xml" }
 bytes = "1.2"
 futures = "0.3"
 http = "0.2"
 hyper = { version = "0.14.20", features = ["server", "http1", "http2", "tcp", "stream"] }
-lambda_http = { version = "0.7.1", optional = true }
 tls-listener = { version = "0.5.1", features = ["rustls", "hyper-h2"] }
 rustls-pemfile = "1.0.1"
 tokio-rustls = "0.23.4"
+lambda_http = { version = "0.7.1" }
+# There is a breaking change in `lambda_runtime` between `0.7.0` and `0.7.1`,
+# and `lambda_http` depends on `0.7` which by default resolves to `0.7.1` but in our CI
+# we are running `minimal-versions` which downgrades `lambda_runtime` to `0.7.0` and fails to compile
+# because of the breaking change. Here we are forcing it to use `lambda_runtime = 0.7.1`.
+lambda_runtime = { version = "0.7.1" }
 num_cpus = "1.13.1"
 parking_lot = "0.12.1"
 pin-project-lite = "0.2"

rust-runtime/aws-smithy-http-server-python/examples/Makefile

@@ -39,7 +39,7 @@ release: codegen
 	ln -sf $(RELEASE_SHARED_LIBRARY_SRC) $(SHARED_LIBRARY_DST)
 
 run: build
-	python $(CUR_DIR)/pokemon_service.py
+	python3 $(CUR_DIR)/pokemon_service.py
 
 test: build
 	cargo test

rust-runtime/aws-smithy-http-server-python/examples/pokemon_service.py

@@ -12,6 +12,7 @@
 
 from libpokemon_service_server_sdk import App
 from libpokemon_service_server_sdk.tls import TlsConfig  # type: ignore
+from libpokemon_service_server_sdk.aws_lambda import LambdaContext  # type: ignore
 from libpokemon_service_server_sdk.error import ResourceNotFoundException  # type: ignore
 from libpokemon_service_server_sdk.input import (  # type: ignore
     DoNothingInput,
@@ -80,6 +81,10 @@ def value(self) -> int:
 #   https://docs.python.org/3/library/multiprocessing.html#sharing-state-between-processes
 @dataclass
 class Context:
+    # Inject Lambda context if service is running on Lambda
+    # NOTE: All the values that will be injected by the framework should be wrapped with `Optional`
+    lambda_ctx: Optional[LambdaContext] = None
+
     # In our case it simulates an in-memory database containing the description of Pikachu in multiple
     # languages.
     _pokemon_database = {
@@ -156,7 +161,7 @@ async def check_content_type_header(request: Request, next: Next) -> Response:
         logging.debug("Found valid `application/json` content type")
     else:
         logging.warning(
-            f"Invalid content type {content_type}, dumping headers: {request.headers}"
+            f"Invalid content type {content_type}, dumping headers: {request.headers.items()}"
         )
     return await next(request)
 
@@ -197,6 +202,18 @@ def do_nothing(_: DoNothingInput) -> DoNothingOutput:
 def get_pokemon_species(
     input: GetPokemonSpeciesInput, context: Context
 ) -> GetPokemonSpeciesOutput:
+    if context.lambda_ctx is not None:
+        logging.debug(
+            "Lambda Context: %s",
+            dict(
+                request_id=context.lambda_ctx.request_id,
+                deadline=context.lambda_ctx.deadline,
+                invoked_function_arn=context.lambda_ctx.invoked_function_arn,
+                function_name=context.lambda_ctx.env_config.function_name,
+                memory=context.lambda_ctx.env_config.memory,
+                version=context.lambda_ctx.env_config.version,
+            ),
+        )
     context.increment_calls_count()
     flavor_text_entries = context.get_pokemon_description(input.name)
     if flavor_text_entries:

rust-runtime/aws-smithy-http-server-python/src/context.rs

@@ -0,0 +1,126 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//! Python context definition.
+
+use http::Extensions;
+use pyo3::{PyObject, PyResult, Python, ToPyObject};
+
+mod lambda;
+pub mod layer;
+#[cfg(test)]
+mod testing;
+
+/// PyContext is a wrapper for context object provided by the user.
+/// It injects some values (currently only [super::lambda::PyLambdaContext]) that is type-hinted by the user.
+///
+///
+/// PyContext is initialised during the startup, it inspects the provided context object for fields
+/// that are type-hinted to inject some values provided by the framework (see [PyContext::new()]).
+///
+/// After finding fields that needs to be injected, [layer::AddPyContextLayer], a [tower::Layer],
+/// populates request-scoped values from incoming request.
+///
+/// And finally PyContext implements [ToPyObject] (so it can by passed to Python handlers)
+/// that provides [PyObject] provided by the user with the additional values injected by the framework.
+#[derive(Clone)]
+pub struct PyContext {
+    inner: PyObject,
+    // TODO(Refactor): We should ideally keep record of injectable fields in a hashmap like:
+    // `injectable_fields: HashMap<Field, Box<dyn Injectable>>` where `Injectable` provides a method to extract a `PyObject` from a `Request`,
+    // but I couldn't find a way to extract a trait object from a Python object.
+    // We could introduce a registry to keep track of every injectable type but I'm not sure that is the best way to do it,
+    // so until we found a good way to achive that, I didn't want to introduce any abstraction here and
+    // keep it simple because we only have one field that is injectable.
+    lambda_ctx: lambda::PyContextLambda,
+}
+
+impl PyContext {
+    pub fn new(inner: PyObject) -> PyResult<Self> {
+        Ok(Self {
+            lambda_ctx: lambda::PyContextLambda::new(inner.clone())?,
+            inner,
+        })
+    }
+
+    pub fn populate_from_extensions(&self, _ext: &Extensions) {
+        self.lambda_ctx
+            .populate_from_extensions(self.inner.clone(), _ext);
+    }
+}
+
+impl ToPyObject for PyContext {
+    fn to_object(&self, _py: Python<'_>) -> PyObject {
+        self.inner.clone()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use http::Extensions;
+    use pyo3::{prelude::*, py_run};
+
+    use super::testing::get_context;
+
+    #[test]
+    fn py_context() -> PyResult<()> {
+        pyo3::prepare_freethreaded_python();
+
+        let ctx = get_context(
+            r#"
+class Context:
+    foo: int = 0
+    bar: str = 'qux'
+
+ctx = Context()
+ctx.foo = 42
+"#,
+        );
+        Python::with_gil(|py| {
+            py_run!(
+                py,
+                ctx,
+                r#"
+assert ctx.foo == 42
+assert ctx.bar == 'qux'
+# Make some modifications
+ctx.foo += 1
+ctx.bar = 'baz'
+"#
+            );
+        });
+
+        ctx.populate_from_extensions(&Extensions::new());
+
+        Python::with_gil(|py| {
+            py_run!(
+                py,
+                ctx,
+                r#"
+# Make sure we are preserving any modifications
+assert ctx.foo == 43
+assert ctx.bar == 'baz'
+"#
+            );
+        });
+
+        Ok(())
+    }
+
+    #[test]
+    fn works_with_none() -> PyResult<()> {
+        // Users can set context to `None` by explicity or implicitly by not providing a custom context class,
+        // it shouldn't be fail in that case.
+
+        pyo3::prepare_freethreaded_python();
+
+        let ctx = get_context("ctx = None");
+        Python::with_gil(|py| {
+            py_run!(py, ctx, "assert ctx is None");
+        });
+
+        Ok(())
+    }
+}