forked from smithy-lang/smithy-rs
-
Notifications
You must be signed in to change notification settings - Fork 0
/
PythonServerOperationHandlerGenerator.kt
141 lines (132 loc) · 6.24 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
136
137
138
139
140
141
/*
* 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
import software.amazon.smithy.rust.codegen.server.smithy.generators.ServerOperationHandlerGenerator
import software.amazon.smithy.rust.codegen.server.smithy.generators.protocol.ServerProtocol
/**
* 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,
protocol: ServerProtocol,
private val operations: List<OperationShape>,
) : ServerOperationHandlerGenerator(codegenContext, protocol, operations) {
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(),
)
override fun render(writer: RustWriter) {
super.render(writer)
renderPythonOperationHandlerImpl(writer)
}
private fun renderPythonOperationHandlerImpl(writer: RustWriter) {
for (operation in operations) {
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,
)
}
}