1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167

/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

//! Support for injecting [PyLambdaContext] to [super::PyContext].

use std::collections::HashSet;

use http::Extensions;
use lambda_http::Context as LambdaContext;
use pyo3::{types::PyDict, PyObject, PyResult, Python};

use crate::{lambda::PyLambdaContext, rich_py_err, util::is_optional_of};

#[derive(Clone)]
pub struct PyContextLambda {
    fields: HashSet<String>,
}

impl PyContextLambda {
    pub fn new(ctx: PyObject) -> PyResult<Self> {
        let fields = Python::with_gil(|py| get_lambda_ctx_fields(py, &ctx))?;
        Ok(Self { fields })
    }

    pub fn populate_from_extensions(&self, ctx: PyObject, ext: &Extensions) {
        if self.fields.is_empty() {
            // Return early without acquiring GIL
            return;
        }

        let lambda_ctx = ext
            .get::<LambdaContext>()
            .cloned()
            .map(PyLambdaContext::new);

        Python::with_gil(|py| {
            for field in self.fields.iter() {
                if let Err(err) = ctx.setattr(py, field.as_str(), lambda_ctx.clone()) {
                    tracing::warn!(field = ?field, error = ?rich_py_err(err), "could not inject `LambdaContext` to context")
                }
            }
        });
    }
}

// Inspects the given `PyObject` to detect fields that type-hinted `PyLambdaContext`.
fn get_lambda_ctx_fields(py: Python, ctx: &PyObject) -> PyResult<HashSet<String>> {
    let typing = py.import("typing")?;
    let hints = match typing
        .call_method1("get_type_hints", (ctx,))
        .and_then(|res| res.extract::<&PyDict>())
    {
        Ok(hints) => hints,
        Err(_) => {
            // `get_type_hints` could fail if `ctx` is `None`, which is the default value
            // for the context if user does not provide a custom class.
            // In that case, this is not really an error and we should just return an empty set.
            return Ok(HashSet::new());
        }
    };

    let mut fields = HashSet::new();
    for (key, value) in hints {
        if is_optional_of::<PyLambdaContext>(py, value)? {
            fields.insert(key.to_string());
        }
    }
    Ok(fields)
}

#[cfg(test)]
mod tests {
    use http::Extensions;
    use lambda_http::Context as LambdaContext;
    use pyo3::{prelude::*, py_run};

    use crate::context::testing::{get_context, lambda_ctx};

    #[test]
    fn py_context_with_lambda_context() -> PyResult<()> {
        pyo3::prepare_freethreaded_python();

        let ctx = get_context(
            r#"
class Context:
    foo: int = 0
    bar: str = 'qux'
    lambda_ctx: typing.Optional[LambdaContext]

ctx = Context()
ctx.foo = 42
"#,
        );
        Python::with_gil(|py| {
            py_run!(
                py,
                ctx,
                r#"
assert ctx.foo == 42
assert ctx.bar == 'qux'
assert not hasattr(ctx, 'lambda_ctx')
"#
            );
        });

        ctx.populate_from_extensions(&extensions_with_lambda_ctx(lambda_ctx("my-req-id", "123")));
        Python::with_gil(|py| {
            py_run!(
                py,
                ctx,
                r#"
assert ctx.lambda_ctx.request_id == "my-req-id"
assert ctx.lambda_ctx.deadline == 123
# Make some modifications
ctx.foo += 1
ctx.bar = 'baz'
"#
            );
        });

        // Assume we are getting a new request but that one doesn't have a `LambdaContext`,
        // in that case we should make fields `None` and shouldn't leak the previous `LambdaContext`.
        ctx.populate_from_extensions(&empty_extensions());
        Python::with_gil(|py| {
            py_run!(
                py,
                ctx,
                r#"
assert ctx.lambda_ctx is None
# 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");
        ctx.populate_from_extensions(&extensions_with_lambda_ctx(lambda_ctx("my-req-id", "123")));
        Python::with_gil(|py| {
            py_run!(py, ctx, "assert ctx is None");
        });

        Ok(())
    }

    fn extensions_with_lambda_ctx(ctx: LambdaContext) -> Extensions {
        let mut exts = empty_extensions();
        exts.insert(ctx);
        exts
    }

    fn empty_extensions() -> Extensions {
        Extensions::new()
    }
}