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

aws_smithy_http_server_python/
logging.rs
