Support changing log levels and setting them in Python (#982)

This introduces a new function `icechunk.set_logs_filter(str)` that
allows setting and changing the logs Icechunk will produce.

On module import, we still initialize logs to the value of ICECHUNK_LOG
environment variable (or error level if not present).

Author: paraseba

docs/docs/icechunk-python/faq.md

@@ -7,3 +7,12 @@ Icechunk is different from normal Zarr stores because it is stateful. In a distr
 **Does `icechunk-python` include logging?**
 
 Yes! Set the environment variable `ICECHUNK_LOG=icechunk=debug` to print debug logs to stdout. Available "levels" in order of increasing verbosity are `error`, `warn`, `info`, `debug`, `trace`. The default level is `error`. The Rust library uses `tracing-subscriber` crate. The `ICECHUNK_LOG` variable can be used to filter logging following that crate's [documentation](https://docs.rs/tracing-subscriber/latest/tracing_subscriber/filter/struct.EnvFilter.html#directives). For example, `ICECHUNK_LOG=trace` will set both icechunk and it's dependencies' log levels to `trace` while `ICECHUNK_LOG=icechunk=trace` will enable the `trace` level for icechunk only. For more complex control `ICECHUNK_LOG=debug,icechunk=trace,rustls=info,h2=info,hyper=info` will set `trace` for `icechunk`, `info` for `rustls`,`hyper`, and `h2` crates, and `debug` for every other crate.
+
+You can also use Python's `os.environ` to set or change the value of the variable. If you change the environment variable after `icechunk` was
+imported, you will need to call `icechunk.set_logs_filter(None)` for changes to take effect.
+
+This function also accepts the filter directive. If you prefer not to use environment variables, you can do:
+
+```python
+icechunk.set_logs_filter("debug,icechunk=trace")
+```

icechunk-python/python/icechunk/__init__.py

@@ -43,6 +43,7 @@
     VirtualChunkSpec,
     __version__,
     initialize_logs,
+    set_logs_filter,
     spec_version,
 )
 from icechunk.credentials import (
@@ -153,6 +154,7 @@
     "s3_static_credentials",
     "s3_storage",
     "s3_store",
+    "set_logs_filter",
     "spec_version",
     "tigris_storage",
 ]

icechunk-python/python/icechunk/_icechunk_python.pyi

@@ -1845,7 +1845,29 @@ def initialize_logs() -> None:
     """
     Initialize the logging system for the library.
 
-    This should be called before any other Icechunk functions are called.
+    Reads the value of the environment variable ICECHUNK_LOG to obtain the filters.
+    This is autamtically called on `import icechunk`.
+    """
+    ...
+
+def set_logs_filter(log_filter_directive: str | None) -> None:
+    """
+    Set filters and log levels for the different modules.
+
+    Examples:
+      - set_logs_filter("trace")  # trace level for all modules
+      - set_logs_filter("error")  # error level for all modules
+      - set_logs_filter("icechunk=debug,info")  # debug level for icechunk, info for everything else
+
+    Full spec for the log_filter_directive syntax is documented in
+    https://docs.rs/tracing-subscriber/latest/tracing_subscriber/filter/struct.EnvFilter.html#directives
+
+    Parameters
+    ----------
+    log_filter_directive: str | None
+        The comma separated list of directives for modules and log levels.
+        If None, the directive will be read from the environment variable
+        ICECHUNK_LOG
     """
     ...
 

icechunk-python/src/lib.rs

@@ -27,6 +27,7 @@ use conflicts::{
 use errors::{IcechunkError, PyConflictError, PyRebaseFailedError};
 use icechunk::{format::format_constants::SpecVersionBin, initialize_tracing};
 use pyo3::prelude::*;
+use pyo3::types::PyMapping;
 use pyo3::wrap_pyfunction;
 use repository::{PyDiff, PyGCSummary, PyManifestFileInfo, PyRepository, PySnapshotInfo};
 use session::PySession;
@@ -69,14 +70,31 @@ fn cli_entrypoint(_py: Python) -> PyResult<()> {
     Ok(())
 }
 
+fn log_filters_from_env(py: Python) -> PyResult<Option<String>> {
+    let os = py.import("os")?;
+    let environ = os.getattr("environ")?;
+    let environ: &Bound<PyMapping> = environ.downcast()?;
+    let value = environ.get_item("ICECHUNK_LOG").ok().and_then(|v| v.extract().ok());
+    Ok(value)
+}
+
 #[pyfunction]
-fn initialize_logs() -> PyResult<()> {
+fn initialize_logs(py: Python) -> PyResult<()> {
     if env::var("ICECHUNK_NO_LOGS").is_err() {
-        initialize_tracing()
+        let log_filter_directive = log_filters_from_env(py)?;
+        initialize_tracing(log_filter_directive.as_deref())
     }
     Ok(())
 }
 
+#[pyfunction]
+fn set_logs_filter(py: Python, log_filter_directive: Option<String>) -> PyResult<()> {
+    let log_filter_directive =
+        log_filter_directive.or_else(|| log_filters_from_env(py).ok().flatten());
+    initialize_tracing(log_filter_directive.as_deref());
+    Ok(())
+}
+
 #[pyfunction]
 /// The spec version that this version of the Icechunk library
 /// uses to write metadata files
@@ -125,6 +143,7 @@ fn _icechunk_python(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyDiff>()?;
     m.add_class::<VirtualChunkSpec>()?;
     m.add_function(wrap_pyfunction!(initialize_logs, m)?)?;
+    m.add_function(wrap_pyfunction!(set_logs_filter, m)?)?;
     m.add_function(wrap_pyfunction!(spec_version, m)?)?;
     m.add_function(wrap_pyfunction!(cli_entrypoint, m)?)?;
 

icechunk-python/tests/test_logs.py

@@ -0,0 +1,57 @@
+import os
+from unittest import mock
+
+import icechunk as ic
+
+
+@mock.patch.dict(os.environ, {"ICECHUNK_LOG": "debug"}, clear=True)
+def test_debug_logs_from_environment(capfd):
+    ic.set_logs_filter(None)
+    ic.Repository.create(storage=ic.in_memory_storage())
+    assert "Creating Repository" in capfd.readouterr().out
+
+
+@mock.patch.dict(os.environ, clear=True)
+def test_no_logs_from_environment(capfd):
+    ic.set_logs_filter(None)
+    ic.Repository.create(storage=ic.in_memory_storage())
+    assert capfd.readouterr().out == ""
+
+
+@mock.patch.dict(os.environ, clear=True)
+def test_change_log_levels_from_env(capfd):
+    # first with logs disabled
+    ic.set_logs_filter(None)
+    ic.Repository.create(storage=ic.in_memory_storage())
+    assert capfd.readouterr().out == ""
+
+    # now with logs enabled
+    with mock.patch.dict(os.environ, {"ICECHUNK_LOG": "debug"}, clear=True):
+        ic.set_logs_filter(None)
+        ic.Repository.create(storage=ic.in_memory_storage())
+        assert "Creating Repository" in capfd.readouterr().out
+
+
+def test_debug_logs_from_argument(capfd):
+    ic.set_logs_filter("debug")
+    ic.Repository.create(storage=ic.in_memory_storage())
+    assert "Creating Repository" in capfd.readouterr().out
+
+
+@mock.patch.dict(os.environ, {"ICECHUNK_LOG": "debug"}, clear=True)
+def test_no_logs_from_argument(capfd):
+    ic.set_logs_filter("false")
+    ic.Repository.create(storage=ic.in_memory_storage())
+    assert capfd.readouterr().out == ""
+
+
+def test_change_log_levels_from_argument(capfd):
+    # first with logs disabled
+    ic.set_logs_filter("")
+    ic.Repository.create(storage=ic.in_memory_storage())
+    assert capfd.readouterr().out == ""
+
+    # now with logs enabled
+    ic.set_logs_filter("debug")
+    ic.Repository.create(storage=ic.in_memory_storage())
+    assert "Creating Repository" in capfd.readouterr().out

icechunk/src/lib.rs

@@ -49,24 +49,60 @@ mod private {
 }
 
 #[cfg(feature = "logs")]
-pub fn initialize_tracing() {
+#[allow(clippy::type_complexity)]
+static LOG_FILTER: std::sync::LazyLock<
+    std::sync::Mutex<
+        Option<
+            tracing_subscriber::reload::Handle<
+                tracing_subscriber::EnvFilter,
+                tracing_subscriber::layer::Layered<
+                    tracing_error::ErrorLayer<tracing_subscriber::Registry>,
+                    tracing_subscriber::Registry,
+                >,
+            >,
+        >,
+    >,
+> = std::sync::LazyLock::new(|| std::sync::Mutex::new(None));
+
+#[cfg(feature = "logs")]
+pub fn initialize_tracing(log_filter_directive: Option<&str>) {
     use tracing_error::ErrorLayer;
     use tracing_subscriber::{
-        EnvFilter, Layer, Registry, layer::SubscriberExt, util::SubscriberInitExt,
+        EnvFilter, Layer, Registry, layer::SubscriberExt, reload, util::SubscriberInitExt,
     };
 
     // We have two Layers. One keeps track of the spans to feed the ICError instances.
     // The other is the one spitting logs to stdout. Filtering only applies to the second Layer.
 
-    let stdout_layer = tracing_subscriber::fmt::layer()
-        .pretty()
-        .with_filter(EnvFilter::from_env("ICECHUNK_LOG"));
+    let filter = log_filter_directive
+        .map(EnvFilter::new)
+        .unwrap_or_else(|| EnvFilter::from_env("ICECHUNK_LOG"));
+    match LOG_FILTER.lock() {
+        Ok(mut guard) => match guard.as_ref() {
+            Some(handle) => {
+                if let Err(err) = handle.reload(filter) {
+                    println!("Error reloading log settings: {}", err)
+                }
+            }
+            None => {
+                let (filter, handle) = reload::Layer::new(filter);
+                *guard = Some(handle);
+                let stdout_layer =
+                    tracing_subscriber::fmt::layer().pretty().with_filter(filter);
 
-    let error_span_layer = ErrorLayer::default();
+                let error_span_layer = ErrorLayer::default();
 
-    if let Err(err) =
-        Registry::default().with(error_span_layer).with(stdout_layer).try_init()
-    {
-        println!("Warning: {}", err);
+                if let Err(err) = Registry::default()
+                    .with(error_span_layer)
+                    .with(stdout_layer)
+                    .try_init()
+                {
+                    println!("Error initializing logs: {}", err);
+                }
+            }
+        },
+        Err(err) => {
+            println!("Error setting up logs: {}", err)
+        }
     }
 }