Add test helper to setup tracing (#12741)

Cargo.lock

@@ -2154,6 +2154,8 @@ dependencies = [
  "salsa",
  "tempfile",
  "tracing",
+ "tracing-subscriber",
+ "tracing-tree",
  "web-time",
  "zip",
 ]

crates/red_knot/docs/tracing.md

@@ -72,6 +72,26 @@ runs or when restoring from a persistent cache. This can be confusing for users
 don't understand why a specific lint violation isn't raised. Instead, change your
 query to return the failure as part of the query's result or use a Salsa accumulator.
 
+## Tracing in tests
+
+You can use `ruff_db::testing::setup_logging` or `ruff_db::testing::setup_logging_with_filter` to set up logging in tests.
+
+```rust
+use ruff_db::testing::setup_logging;
+
+#[test]
+fn test() {
+  let _logging = setup_logging();
+
+  tracing::info!("This message will be printed to stderr");
+}
+```
+
+Note: Most test runners capture stderr and only show its output when a test fails.
+
+Note also that `setup_logging` only sets up logging for the current thread because [`set_global_default`](https://docs.rs/tracing/latest/tracing/subscriber/fn.set_global_default.html) can only be
+called **once**.
+
 ## Release builds
 
 `trace!` events are removed in release builds.

crates/red_knot_python_semantic/Cargo.toml

@@ -34,13 +34,14 @@ walkdir = { workspace = true }
 zip = { workspace = true, features = ["zstd", "deflate"] }
 
 [dev-dependencies]
+ruff_db = { workspace = true, features = ["os", "testing"]}
+ruff_python_parser = { workspace = true }
+
 anyhow = { workspace = true }
 insta = { workspace = true }
 tempfile = { workspace = true }
 walkdir = { workspace = true }
 zip = { workspace = true }
-ruff_python_parser = { workspace = true }
-ruff_db = { workspace = true, features = ["os"]}
 
 [lints]
 workspace = true

crates/red_knot_workspace/src/db.rs

@@ -129,7 +129,18 @@ impl SourceDb for RootDatabase {
 
 #[salsa::db]
 impl salsa::Database for RootDatabase {
-    fn salsa_event(&self, _event: &dyn Fn() -> Event) {}
+    fn salsa_event(&self, event: &dyn Fn() -> Event) {
+        if !tracing::enabled!(tracing::Level::TRACE) {
+            return;
+        }
+
+        let event = event();
+        if matches!(event.kind, salsa::EventKind::WillCheckCancellation { .. }) {
+            return;
+        }
+
+        tracing::trace!("Salsa event: {event:?}");
+    }
 }
 
 #[salsa::db]

crates/ruff_db/Cargo.toml

@@ -28,6 +28,8 @@ matchit = { workspace = true }
 salsa = { workspace = true }
 path-slash = { workspace = true }
 tracing = { workspace = true }
+tracing-subscriber = { workspace = true, optional = true }
+tracing-tree = { workspace = true, optional = true }
 rustc-hash = { workspace = true }
 
 [target.'cfg(not(target_arch="wasm32"))'.dependencies]
@@ -44,3 +46,5 @@ tempfile = { workspace = true }
 [features]
 cache = ["ruff_cache"]
 os = ["ignore"]
+# Exposes testing utilities.
+testing = ["tracing-subscriber", "tracing-tree"]

crates/ruff_db/src/lib.rs

@@ -12,6 +12,7 @@ pub mod parsed;
 pub mod program;
 pub mod source;
 pub mod system;
+#[cfg(feature = "testing")]
 pub mod testing;
 pub mod vendored;
 

crates/ruff_db/src/testing.rs

@@ -1,5 +1,8 @@
 //! Test helpers for working with Salsa databases
 
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::EnvFilter;
+
 pub fn assert_function_query_was_not_run<Db, Q, QDb, I, R>(
     db: &Db,
     query: Q,
@@ -94,6 +97,116 @@ fn query_name<Q>(_query: &Q) -> &'static str {
         .unwrap_or(full_qualified_query_name)
 }
 
+/// Sets up logging for the current thread. It captures all `red_knot` and `ruff` events.
+///
+/// Useful for capturing the tracing output in a failing test.
+///
+/// # Examples
+/// ```
+/// use ruff_db::testing::setup_logging;
+/// let _logging = setup_logging();
+///
+/// tracing::info!("This message will be printed to stderr");
+/// ```
+pub fn setup_logging() -> LoggingGuard {
+    LoggingBuilder::new().build()
+}
+
+/// Sets up logging for the current thread and uses the passed filter to filter the shown events.
+/// Useful for capturing the tracing output in a failing test.
+///
+/// # Examples
+/// ```
+/// use ruff_db::testing::setup_logging_with_filter;
+/// let _logging = setup_logging_with_filter("red_knot_module_resolver::resolver");
+/// ```
+///
+/// # Filter
+/// See [`tracing_subscriber::EnvFilter`] for the `filter`'s syntax.
+///
+pub fn setup_logging_with_filter(filter: &str) -> Option<LoggingGuard> {
+    LoggingBuilder::with_filter(filter).map(LoggingBuilder::build)
+}
+
+#[derive(Debug)]
+pub struct LoggingBuilder {
+    filter: EnvFilter,
+    hierarchical: bool,
+}
+
+impl LoggingBuilder {
+    pub fn new() -> Self {
+        Self {
+            filter: EnvFilter::default()
+                .add_directive(
+                    "red_knot=trace"
+                        .parse()
+                        .expect("Hardcoded directive to be valid"),
+                )
+                .add_directive(
+                    "ruff=trace"
+                        .parse()
+                        .expect("Hardcoded directive to be valid"),
+                ),
+            hierarchical: true,
+        }
+    }
+
+    pub fn with_filter(filter: &str) -> Option<Self> {
+        let filter = EnvFilter::builder().parse(filter).ok()?;
+
+        Some(Self {
+            filter,
+            hierarchical: true,
+        })
+    }
+
+    pub fn with_hierarchical(mut self, hierarchical: bool) -> Self {
+        self.hierarchical = hierarchical;
+        self
+    }
+
+    pub fn build(self) -> LoggingGuard {
+        let registry = tracing_subscriber::registry().with(self.filter);
+
+        let guard = if self.hierarchical {
+            let subscriber = registry.with(
+                tracing_tree::HierarchicalLayer::default()
+                    .with_indent_lines(true)
+                    .with_indent_amount(2)
+                    .with_bracketed_fields(true)
+                    .with_thread_ids(true)
+                    .with_targets(true)
+                    .with_writer(std::io::stderr)
+                    .with_timer(tracing_tree::time::Uptime::default()),
+            );
+
+            tracing::subscriber::set_default(subscriber)
+        } else {
+            let subscriber = registry.with(
+                tracing_subscriber::fmt::layer()
+                    .compact()
+                    .with_writer(std::io::stderr)
+                    .with_timer(tracing_subscriber::fmt::time()),
+            );
+
+            tracing::subscriber::set_default(subscriber)
+        };
+
+        LoggingGuard { _guard: guard }
+    }
+}
+
+impl Default for LoggingBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+pub struct LoggingGuard {
+    _guard: tracing::subscriber::DefaultGuard,
+}
+
 #[test]
 fn query_was_not_run() {
     use crate::tests::TestDb;