-
Notifications
You must be signed in to change notification settings - Fork 180
/
PythonServerOperationHandlerGenerator.kt
135 lines (126 loc) · 5.75 KB
/
PythonServerOperationHandlerGenerator.kt
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
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/
package software.amazon.smithy.rust.codegen.server.python.smithy.generators
import software.amazon.smithy.model.shapes.OperationShape
import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
import software.amazon.smithy.rust.codegen.core.rustlang.Writable
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.CodegenContext
import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
import software.amazon.smithy.rust.codegen.server.python.smithy.PythonServerCargoDependency
/**
* The Rust code responsible to run the Python business logic on the Python interpreter
* is implemented in this class, which inherits from [ServerOperationHandlerGenerator].
*
* We codegenerate all operations handlers (steps usually left to the developer in a pure
* Rust application), which are built into a `Router` by [PythonApplicationGenerator].
*
* To call a Python function from Rust, anything dealing with Python runs inside an async
* block that allows to catch stack traces. The handler function is extracted from `PyHandler`
* and called with the necessary arguments inside a blocking Tokio task.
* At the end the block is awaited and errors are collected and reported.
*
* To call a Python coroutine, the same happens, but scheduled in a `tokio::Future`.
*/
class PythonServerOperationHandlerGenerator(
codegenContext: CodegenContext,
private val operation: OperationShape,
) {
private val symbolProvider = codegenContext.symbolProvider
private val runtimeConfig = codegenContext.runtimeConfig
private val codegenScope =
arrayOf(
"SmithyPython" to PythonServerCargoDependency.smithyHttpServerPython(runtimeConfig).toType(),
"SmithyServer" to PythonServerCargoDependency.smithyHttpServer(runtimeConfig).toType(),
"pyo3" to PythonServerCargoDependency.PyO3.toType(),
"pyo3_asyncio" to PythonServerCargoDependency.PyO3Asyncio.toType(),
"tokio" to PythonServerCargoDependency.Tokio.toType(),
"tracing" to PythonServerCargoDependency.Tracing.toType(),
)
fun render(writer: RustWriter) {
renderPythonOperationHandlerImpl(writer)
}
private fun renderPythonOperationHandlerImpl(writer: RustWriter) {
val operationName = symbolProvider.toSymbol(operation).name
val input = "crate::input::${operationName}Input"
val output = "crate::output::${operationName}Output"
val error = "crate::error::${operationName}Error"
val fnName = operationName.toSnakeCase()
writer.rustTemplate(
"""
/// Python handler for operation `$operationName`.
pub(crate) async fn $fnName(
input: $input,
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.
let result = if handler.is_coroutine {
#{PyCoroutine:W}
} else {
#{PyFunction:W}
};
#{PyError:W}
}
""",
*codegenScope,
"PyCoroutine" to renderPyCoroutine(fnName, output),
"PyFunction" to renderPyFunction(fnName, output),
"PyError" to renderPyError(),
)
}
private fun renderPyFunction(name: String, output: String): Writable =
writable {
rustTemplate(
"""
#{tracing}::trace!(name = "$name", "executing python handler function");
#{pyo3}::Python::with_gil(|py| {
let pyhandler: &#{pyo3}::types::PyFunction = handler.extract(py)?;
let output = if handler.args == 1 {
pyhandler.call1((input,))?
} else {
pyhandler.call1((input, #{pyo3}::ToPyObject::to_object(&state.0, py)))?
};
output.extract::<$output>()
})
""",
*codegenScope,
)
}
private fun renderPyCoroutine(name: String, output: String): Writable =
writable {
rustTemplate(
"""
#{tracing}::trace!(name = "$name", "executing python handler coroutine");
let result = #{pyo3}::Python::with_gil(|py| {
let pyhandler: &#{pyo3}::types::PyFunction = handler.extract(py)?;
let coroutine = if handler.args == 1 {
pyhandler.call1((input,))?
} else {
pyhandler.call1((input, #{pyo3}::ToPyObject::to_object(&state.0, py)))?
};
#{pyo3_asyncio}::tokio::into_future(coroutine)
})?;
result.await.map(|r| #{pyo3}::Python::with_gil(|py| r.extract::<$output>(py)))?
""",
*codegenScope,
)
}
private fun renderPyError(): Writable =
writable {
rustTemplate(
"""
// Catch and record a Python traceback.
result.map_err(|e| {
let rich_py_err = #{SmithyPython}::rich_py_err(#{pyo3}::Python::with_gil(|py| { e.clone_ref(py) }));
#{tracing}::error!(error = ?rich_py_err, "handler error");
e.into()
})
""",
*codegenScope,
)
}
}