Skip to content

Commit

Permalink
Python: Allow injecting Lambda Context (#1985)
Browse files Browse the repository at this point in the history
* 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 <[email protected]>

* 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 <[email protected]>
  • Loading branch information
unexge and crisidev authored Dec 12, 2022
1 parent 9d2d088 commit 3fb9096
Show file tree
Hide file tree
Showing 14 changed files with 896 additions and 49 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -131,6 +115,4 @@ val DECORATORS = listOf(
PubUsePythonTypesDecorator(),
// Render the Python shared library export.
PythonExportModuleDecorator(),
// Add the `aws-lambda` feature flag
PythonFeatureFlagsDecorator(),
)
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ class PythonServerModuleGenerator(
renderPyLogging()
renderPyMiddlewareTypes()
renderPyTlsTypes()
renderPyLambdaTypes()
renderPyApplicationType()
}
}
Expand Down Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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>()
})
Expand All @@ -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)
})?;
Expand Down
12 changes: 7 additions & 5 deletions rust-runtime/aws-smithy-http-server-python/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 = {
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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:
Expand Down
126 changes: 126 additions & 0 deletions rust-runtime/aws-smithy-http-server-python/src/context.rs
Original file line number Diff line number Diff line change
@@ -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(())
}
}
Loading

0 comments on commit 3fb9096

Please sign in to comment.