aws_smithy_http_server_python/
logging.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Rust `tracing` and Python `logging` setup and utilities.
7
8use std::{path::PathBuf, str::FromStr};
9
10use pyo3::prelude::*;
11#[cfg(not(test))]
12use tracing::span;
13use tracing::Level;
14use tracing_appender::non_blocking::WorkerGuard;
15use tracing_subscriber::{
16    fmt::{self, writer::MakeWriterExt},
17    layer::SubscriberExt,
18    util::SubscriberInitExt,
19    Layer,
20};
21
22use crate::error::PyException;
23
24#[derive(Debug, Default)]
25enum Format {
26    Json,
27    Pretty,
28    #[default]
29    Compact,
30}
31
32#[derive(Debug, PartialEq, Eq)]
33struct InvalidFormatError;
34
35impl FromStr for Format {
36    type Err = InvalidFormatError;
37
38    fn from_str(s: &str) -> Result<Self, Self::Err> {
39        match s {
40            "pretty" => Ok(Self::Pretty),
41            "json" => Ok(Self::Json),
42            "compact" => Ok(Self::Compact),
43            _ => Err(InvalidFormatError),
44        }
45    }
46}
47
48/// Setup tracing-subscriber to log on console or to a hourly rolling file.
49fn setup_tracing_subscriber(
50    level: Option<u8>,
51    logfile: Option<PathBuf>,
52    format: Option<String>,
53) -> PyResult<Option<WorkerGuard>> {
54    let format = match format {
55        Some(format) => Format::from_str(&format).unwrap_or_else(|_| {
56            tracing::error!("unknown format '{format}', falling back to default formatter");
57            Format::default()
58        }),
59        None => Format::default(),
60    };
61
62    let appender = match logfile {
63        Some(logfile) => {
64            let parent = logfile.parent().ok_or_else(|| {
65                PyException::new_err(format!(
66                    "Tracing setup failed: unable to extract dirname from path {}",
67                    logfile.display()
68                ))
69            })?;
70            let filename = logfile.file_name().ok_or_else(|| {
71                PyException::new_err(format!(
72                    "Tracing setup failed: unable to extract basename from path {}",
73                    logfile.display()
74                ))
75            })?;
76            let file_appender = tracing_appender::rolling::hourly(parent, filename);
77            let (appender, guard) = tracing_appender::non_blocking(file_appender);
78            Some((appender, guard))
79        }
80        None => None,
81    };
82
83    let tracing_level = match level {
84        Some(40u8) => Level::ERROR,
85        Some(30u8) => Level::WARN,
86        Some(20u8) => Level::INFO,
87        Some(10u8) => Level::DEBUG,
88        None => Level::INFO,
89        _ => Level::TRACE,
90    };
91
92    let formatter = fmt::Layer::new().with_line_number(true).with_level(true);
93
94    match appender {
95        Some((appender, guard)) => {
96            let formatter = formatter.with_writer(appender.with_max_level(tracing_level));
97            let formatter = match format {
98                Format::Json => formatter.json().boxed(),
99                Format::Compact => formatter.compact().boxed(),
100                Format::Pretty => formatter.pretty().boxed(),
101            };
102            tracing_subscriber::registry().with(formatter).init();
103            Ok(Some(guard))
104        }
105        None => {
106            let formatter = formatter.with_writer(std::io::stdout.with_max_level(tracing_level));
107            let formatter = match format {
108                Format::Json => formatter.json().boxed(),
109                Format::Compact => formatter.compact().boxed(),
110                Format::Pretty => formatter.pretty().boxed(),
111            };
112            tracing_subscriber::registry().with(formatter).init();
113            Ok(None)
114        }
115    }
116}
117
118/// Modifies the Python `logging` module to deliver its log messages using [tracing::Subscriber] events.
119///
120/// To achieve this goal, the following changes are made to the module:
121/// - A new builtin function `logging.py_tracing_event` transcodes `logging.LogRecord`s to `tracing::Event`s. This function
122///   is not exported in `logging.__all__`, as it is not intended to be called directly.
123/// - A new class `logging.TracingHandler` provides a `logging.Handler` that delivers all records to `python_tracing`.
124///
125/// :param level typing.Optional\[int\]:
126/// :param logfile typing.Optional\[pathlib.Path\]:
127/// :param format typing.Optional\[typing.Literal\['compact', 'pretty', 'json'\]\]:
128/// :rtype None:
129#[pyclass(name = "TracingHandler")]
130#[derive(Debug)]
131pub struct PyTracingHandler {
132    _guard: Option<WorkerGuard>,
133}
134
135#[pymethods]
136impl PyTracingHandler {
137    #[pyo3(text_signature = "($self, level=None, logfile=None, format=None)")]
138    #[new]
139    fn newpy(
140        py: Python,
141        level: Option<u8>,
142        logfile: Option<PathBuf>,
143        format: Option<String>,
144    ) -> PyResult<Self> {
145        let _guard = setup_tracing_subscriber(level, logfile, format)?;
146        let logging = py.import("logging")?;
147        let root = logging.getattr("root")?;
148        root.setattr("level", level)?;
149        // TODO(Investigate why the file appender just create the file and does not write anything, event after holding the guard)
150        Ok(Self { _guard })
151    }
152
153    /// :rtype typing.Any:
154    fn handler(&self, py: Python) -> PyResult<Py<PyAny>> {
155        let logging = py.import("logging")?;
156        logging.setattr(
157            "py_tracing_event",
158            wrap_pyfunction!(py_tracing_event, logging)?,
159        )?;
160
161        let pycode = r#"
162class TracingHandler(Handler):
163    """ Python logging to Rust tracing handler. """
164    def emit(self, record):
165        py_tracing_event(
166            record.levelno, record.getMessage(), record.module,
167            record.filename, record.lineno, record.process
168        )
169"#;
170        py.run(pycode, Some(logging.dict()), None)?;
171        let all = logging.index()?;
172        all.append("TracingHandler")?;
173        let handler = logging.getattr("TracingHandler")?;
174        Ok(handler.call0()?.into_py(py))
175    }
176}
177
178/// Consumes a Python `logging.LogRecord` and emits a Rust [tracing::Event] instead.
179#[cfg(not(test))]
180#[pyfunction]
181#[pyo3(text_signature = "(level, record, message, module, filename, line, pid)")]
182pub fn py_tracing_event(
183    level: u8,
184    message: &str,
185    module: &str,
186    filename: &str,
187    lineno: usize,
188    pid: usize,
189) -> PyResult<()> {
190    let span = span!(
191        Level::TRACE,
192        "python",
193        pid = pid,
194        module = module,
195        filename = filename,
196        lineno = lineno
197    );
198    let _guard = span.enter();
199    match level {
200        40 => tracing::error!("{message}"),
201        30 => tracing::warn!("{message}"),
202        20 => tracing::info!("{message}"),
203        10 => tracing::debug!("{message}"),
204        _ => tracing::trace!("{message}"),
205    };
206    Ok(())
207}
208
209#[cfg(test)]
210#[pyfunction]
211#[pyo3(text_signature = "(level, record, message, module, filename, line, pid)")]
212pub fn py_tracing_event(
213    _level: u8,
214    message: &str,
215    _module: &str,
216    _filename: &str,
217    _line: usize,
218    _pid: usize,
219) -> PyResult<()> {
220    pretty_assertions::assert_eq!(message.to_string(), "a message");
221    Ok(())
222}
223
224#[cfg(test)]
225mod tests {
226    use pyo3::types::PyDict;
227
228    use super::*;
229
230    #[test]
231    fn tracing_handler_is_injected_in_python() {
232        crate::tests::initialize();
233        Python::with_gil(|py| {
234            let handler = PyTracingHandler::newpy(py, Some(10), None, None).unwrap();
235            let kwargs = PyDict::new(py);
236            kwargs
237                .set_item("handlers", vec![handler.handler(py).unwrap()])
238                .unwrap();
239            let logging = py.import("logging").unwrap();
240            let basic_config = logging.getattr("basicConfig").unwrap();
241            basic_config.call((), Some(kwargs)).unwrap();
242            logging.call_method1("info", ("a message",)).unwrap();
243        });
244    }
245}
aws_smithy_http_server_python/logging.rs

aws_smithy_http_server_python/
logging.rs
