diff --git a/Cargo.lock b/Cargo.lock index 70aab8897b13..37d53be81157 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -280,6 +280,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "clap" version = "4.5.16" @@ -339,7 +345,7 @@ dependencies = [ ] [[package]] -name = "conn-pool" +name = "conn_pool" version = "0.1.0" dependencies = [ "anyhow", @@ -537,6 +543,18 @@ dependencies = [ "thiserror", ] +[[package]] +name = "edbrust" +version = "0.1.0" +dependencies = [ + "conn_pool", + "http 0.1.0", + "pgrust", + "pyo3", + "pyo3_util", + "tokio", +] + [[package]] name = "edgedb-errors" version = "0.4.2" @@ -1341,6 +1359,18 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "nom" version = "7.1.3" @@ -1741,6 +1771,19 @@ dependencies = [ "syn", ] +[[package]] +name = "pyo3_util" +version = "0.1.0" +dependencies = [ + "nix", + "pyo3", + "scopeguard", + "thiserror", + "tokio", + "tracing", + "tracing-subscriber", +] + [[package]] name = "quote" version = "1.0.37" diff --git a/Cargo.toml b/Cargo.toml index d5bb17fc60d2..411980c916ba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,17 +4,23 @@ members = [ "edb/edgeql-parser/edgeql-parser-derive", "edb/edgeql-parser/edgeql-parser-python", "edb/graphql-rewrite", - "edb/server/conn_pool", - "edb/server/pgrust", - "edb/server/http", + "edb/native/conn_pool", + "edb/native/pgrust", + "edb/native/http", + "edb/native/pyo3_util", + "edb/server/_rust_native" ] resolver = "2" [workspace.dependencies] -pyo3 = { version = "0.23.1", features = ["extension-module", "serde"] } +pyo3 = { version = "0.23.1", features = ["extension-module", "serde", "macros"] } tokio = { version = "1", features = ["rt", "rt-multi-thread", "macros", "time", "sync", "net", "io-util"] } tracing = "0.1.40" tracing-subscriber = "0.3.18" +conn_pool = { path = "edb/native/conn_pool" } +pgrust = { path = "edb/native/pgrust" } +http = { path = "edb/native/http" } +pyo3_util = { path = "edb/native/pyo3_util" } [profile.release] debug = true diff --git a/edb/server/conn_pool/Cargo.toml b/edb/native/conn_pool/Cargo.toml similarity index 98% rename from edb/server/conn_pool/Cargo.toml rename to edb/native/conn_pool/Cargo.toml index 10a5966eae59..f9a649cb614a 100644 --- a/edb/server/conn_pool/Cargo.toml +++ b/edb/native/conn_pool/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "conn-pool" +name = "conn_pool" version = "0.1.0" license = "MIT/Apache-2.0" authors = ["MagicStack Inc. "] diff --git a/edb/server/conn_pool/README.md b/edb/native/conn_pool/README.md similarity index 100% rename from edb/server/conn_pool/README.md rename to edb/native/conn_pool/README.md diff --git a/edb/server/conn_pool/src/algo.rs b/edb/native/conn_pool/src/algo.rs similarity index 100% rename from edb/server/conn_pool/src/algo.rs rename to edb/native/conn_pool/src/algo.rs diff --git a/edb/server/conn_pool/src/bin/optimizer.rs b/edb/native/conn_pool/src/bin/optimizer.rs similarity index 100% rename from edb/server/conn_pool/src/bin/optimizer.rs rename to edb/native/conn_pool/src/bin/optimizer.rs diff --git a/edb/server/conn_pool/src/block.rs b/edb/native/conn_pool/src/block.rs similarity index 100% rename from edb/server/conn_pool/src/block.rs rename to edb/native/conn_pool/src/block.rs diff --git a/edb/server/conn_pool/src/conn.rs b/edb/native/conn_pool/src/conn.rs similarity index 100% rename from edb/server/conn_pool/src/conn.rs rename to edb/native/conn_pool/src/conn.rs diff --git a/edb/server/conn_pool/src/drain.rs b/edb/native/conn_pool/src/drain.rs similarity index 100% rename from edb/server/conn_pool/src/drain.rs rename to edb/native/conn_pool/src/drain.rs diff --git a/edb/server/conn_pool/src/lib.rs b/edb/native/conn_pool/src/lib.rs similarity index 97% rename from edb/server/conn_pool/src/lib.rs rename to edb/native/conn_pool/src/lib.rs index 7ec8d3ed51ff..a18e0ed67cf7 100644 --- a/edb/server/conn_pool/src/lib.rs +++ b/edb/native/conn_pool/src/lib.rs @@ -25,4 +25,4 @@ pub use pool::{Pool, PoolConfig, PoolHandle}; pub mod test; #[cfg(feature = "python_extension")] -mod python; +pub mod python; diff --git a/edb/server/conn_pool/src/metrics.rs b/edb/native/conn_pool/src/metrics.rs similarity index 100% rename from edb/server/conn_pool/src/metrics.rs rename to edb/native/conn_pool/src/metrics.rs diff --git a/edb/server/conn_pool/src/pool.rs b/edb/native/conn_pool/src/pool.rs similarity index 100% rename from edb/server/conn_pool/src/pool.rs rename to edb/native/conn_pool/src/pool.rs diff --git a/edb/server/conn_pool/src/python.rs b/edb/native/conn_pool/src/python.rs similarity index 99% rename from edb/server/conn_pool/src/python.rs rename to edb/native/conn_pool/src/python.rs index 2d05b73e25f8..c6f7c4b08d7e 100644 --- a/edb/server/conn_pool/src/python.rs +++ b/edb/native/conn_pool/src/python.rs @@ -500,7 +500,7 @@ impl LoggingGuard { } #[pymodule] -fn _conn_pool(py: Python, m: &Bound) -> PyResult<()> { +pub fn _conn_pool(py: Python, m: &Bound) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add("InternalError", py.get_type::())?; diff --git a/edb/server/conn_pool/src/test/mod.rs b/edb/native/conn_pool/src/test/mod.rs similarity index 100% rename from edb/server/conn_pool/src/test/mod.rs rename to edb/native/conn_pool/src/test/mod.rs diff --git a/edb/server/conn_pool/src/test/spec.rs b/edb/native/conn_pool/src/test/spec.rs similarity index 100% rename from edb/server/conn_pool/src/test/spec.rs rename to edb/native/conn_pool/src/test/spec.rs diff --git a/edb/server/conn_pool/src/waitqueue.rs b/edb/native/conn_pool/src/waitqueue.rs similarity index 100% rename from edb/server/conn_pool/src/waitqueue.rs rename to edb/native/conn_pool/src/waitqueue.rs diff --git a/edb/server/http/Cargo.toml b/edb/native/http/Cargo.toml similarity index 100% rename from edb/server/http/Cargo.toml rename to edb/native/http/Cargo.toml diff --git a/edb/server/http/src/lib.rs b/edb/native/http/src/lib.rs similarity index 69% rename from edb/server/http/src/lib.rs rename to edb/native/http/src/lib.rs index 83610f1accaa..71429ec8ff93 100644 --- a/edb/server/http/src/lib.rs +++ b/edb/native/http/src/lib.rs @@ -1,2 +1,2 @@ #[cfg(feature = "python_extension")] -mod python; +pub mod python; diff --git a/edb/server/http/src/python.rs b/edb/native/http/src/python.rs similarity index 99% rename from edb/server/http/src/python.rs rename to edb/native/http/src/python.rs index 83e113afd76a..595ca88e6212 100644 --- a/edb/server/http/src/python.rs +++ b/edb/native/http/src/python.rs @@ -664,7 +664,7 @@ impl Http { } #[pymodule] -fn _http(py: Python, m: &Bound) -> PyResult<()> { +pub fn _http(py: Python, m: &Bound) -> PyResult<()> { m.add_class::()?; m.add("InternalError", py.get_type::())?; diff --git a/edb/server/pgrust/Cargo.toml b/edb/native/pgrust/Cargo.toml similarity index 100% rename from edb/server/pgrust/Cargo.toml rename to edb/native/pgrust/Cargo.toml diff --git a/edb/server/pgrust/examples/connect.rs b/edb/native/pgrust/examples/connect.rs similarity index 100% rename from edb/server/pgrust/examples/connect.rs rename to edb/native/pgrust/examples/connect.rs diff --git a/edb/server/pgrust/examples/dsn.rs b/edb/native/pgrust/examples/dsn.rs similarity index 100% rename from edb/server/pgrust/examples/dsn.rs rename to edb/native/pgrust/examples/dsn.rs diff --git a/edb/server/pgrust/src/auth/md5.rs b/edb/native/pgrust/src/auth/md5.rs similarity index 100% rename from edb/server/pgrust/src/auth/md5.rs rename to edb/native/pgrust/src/auth/md5.rs diff --git a/edb/server/pgrust/src/auth/mod.rs b/edb/native/pgrust/src/auth/mod.rs similarity index 100% rename from edb/server/pgrust/src/auth/mod.rs rename to edb/native/pgrust/src/auth/mod.rs diff --git a/edb/server/pgrust/src/auth/scram.rs b/edb/native/pgrust/src/auth/scram.rs similarity index 100% rename from edb/server/pgrust/src/auth/scram.rs rename to edb/native/pgrust/src/auth/scram.rs diff --git a/edb/server/pgrust/src/auth/stringprep.rs b/edb/native/pgrust/src/auth/stringprep.rs similarity index 100% rename from edb/server/pgrust/src/auth/stringprep.rs rename to edb/native/pgrust/src/auth/stringprep.rs diff --git a/edb/server/pgrust/src/auth/stringprep_table.rs b/edb/native/pgrust/src/auth/stringprep_table.rs similarity index 100% rename from edb/server/pgrust/src/auth/stringprep_table.rs rename to edb/native/pgrust/src/auth/stringprep_table.rs diff --git a/edb/server/pgrust/src/auth/stringprep_table_prep.py b/edb/native/pgrust/src/auth/stringprep_table_prep.py similarity index 100% rename from edb/server/pgrust/src/auth/stringprep_table_prep.py rename to edb/native/pgrust/src/auth/stringprep_table_prep.py diff --git a/edb/server/pgrust/src/connection/conn.rs b/edb/native/pgrust/src/connection/conn.rs similarity index 100% rename from edb/server/pgrust/src/connection/conn.rs rename to edb/native/pgrust/src/connection/conn.rs diff --git a/edb/server/pgrust/src/connection/dsn/host.rs b/edb/native/pgrust/src/connection/dsn/host.rs similarity index 100% rename from edb/server/pgrust/src/connection/dsn/host.rs rename to edb/native/pgrust/src/connection/dsn/host.rs diff --git a/edb/server/pgrust/src/connection/dsn/mod.rs b/edb/native/pgrust/src/connection/dsn/mod.rs similarity index 100% rename from edb/server/pgrust/src/connection/dsn/mod.rs rename to edb/native/pgrust/src/connection/dsn/mod.rs diff --git a/edb/server/pgrust/src/connection/dsn/params.rs b/edb/native/pgrust/src/connection/dsn/params.rs similarity index 100% rename from edb/server/pgrust/src/connection/dsn/params.rs rename to edb/native/pgrust/src/connection/dsn/params.rs diff --git a/edb/server/pgrust/src/connection/dsn/passfile.rs b/edb/native/pgrust/src/connection/dsn/passfile.rs similarity index 100% rename from edb/server/pgrust/src/connection/dsn/passfile.rs rename to edb/native/pgrust/src/connection/dsn/passfile.rs diff --git a/edb/server/pgrust/src/connection/dsn/raw_params.rs b/edb/native/pgrust/src/connection/dsn/raw_params.rs similarity index 100% rename from edb/server/pgrust/src/connection/dsn/raw_params.rs rename to edb/native/pgrust/src/connection/dsn/raw_params.rs diff --git a/edb/server/pgrust/src/connection/dsn/url.rs b/edb/native/pgrust/src/connection/dsn/url.rs similarity index 100% rename from edb/server/pgrust/src/connection/dsn/url.rs rename to edb/native/pgrust/src/connection/dsn/url.rs diff --git a/edb/server/pgrust/src/connection/mod.rs b/edb/native/pgrust/src/connection/mod.rs similarity index 100% rename from edb/server/pgrust/src/connection/mod.rs rename to edb/native/pgrust/src/connection/mod.rs diff --git a/edb/server/pgrust/src/connection/openssl.rs b/edb/native/pgrust/src/connection/openssl.rs similarity index 100% rename from edb/server/pgrust/src/connection/openssl.rs rename to edb/native/pgrust/src/connection/openssl.rs diff --git a/edb/server/pgrust/src/connection/raw_conn.rs b/edb/native/pgrust/src/connection/raw_conn.rs similarity index 100% rename from edb/server/pgrust/src/connection/raw_conn.rs rename to edb/native/pgrust/src/connection/raw_conn.rs diff --git a/edb/server/pgrust/src/connection/stream.rs b/edb/native/pgrust/src/connection/stream.rs similarity index 100% rename from edb/server/pgrust/src/connection/stream.rs rename to edb/native/pgrust/src/connection/stream.rs diff --git a/edb/server/pgrust/src/connection/tokio.rs b/edb/native/pgrust/src/connection/tokio.rs similarity index 100% rename from edb/server/pgrust/src/connection/tokio.rs rename to edb/native/pgrust/src/connection/tokio.rs diff --git a/edb/server/pgrust/src/errors/mod.rs b/edb/native/pgrust/src/errors/mod.rs similarity index 100% rename from edb/server/pgrust/src/errors/mod.rs rename to edb/native/pgrust/src/errors/mod.rs diff --git a/edb/server/pgrust/src/handshake/client_state_machine.rs b/edb/native/pgrust/src/handshake/client_state_machine.rs similarity index 100% rename from edb/server/pgrust/src/handshake/client_state_machine.rs rename to edb/native/pgrust/src/handshake/client_state_machine.rs diff --git a/edb/server/pgrust/src/handshake/mod.rs b/edb/native/pgrust/src/handshake/mod.rs similarity index 100% rename from edb/server/pgrust/src/handshake/mod.rs rename to edb/native/pgrust/src/handshake/mod.rs diff --git a/edb/server/pgrust/src/handshake/server_state_machine.rs b/edb/native/pgrust/src/handshake/server_state_machine.rs similarity index 100% rename from edb/server/pgrust/src/handshake/server_state_machine.rs rename to edb/native/pgrust/src/handshake/server_state_machine.rs diff --git a/edb/server/pgrust/src/lib.rs b/edb/native/pgrust/src/lib.rs similarity index 100% rename from edb/server/pgrust/src/lib.rs rename to edb/native/pgrust/src/lib.rs diff --git a/edb/server/pgrust/src/protocol/arrays.rs b/edb/native/pgrust/src/protocol/arrays.rs similarity index 100% rename from edb/server/pgrust/src/protocol/arrays.rs rename to edb/native/pgrust/src/protocol/arrays.rs diff --git a/edb/server/pgrust/src/protocol/buffer.rs b/edb/native/pgrust/src/protocol/buffer.rs similarity index 100% rename from edb/server/pgrust/src/protocol/buffer.rs rename to edb/native/pgrust/src/protocol/buffer.rs diff --git a/edb/server/pgrust/src/protocol/datatypes.rs b/edb/native/pgrust/src/protocol/datatypes.rs similarity index 100% rename from edb/server/pgrust/src/protocol/datatypes.rs rename to edb/native/pgrust/src/protocol/datatypes.rs diff --git a/edb/server/pgrust/src/protocol/definition.rs b/edb/native/pgrust/src/protocol/definition.rs similarity index 100% rename from edb/server/pgrust/src/protocol/definition.rs rename to edb/native/pgrust/src/protocol/definition.rs diff --git a/edb/server/pgrust/src/protocol/gen.rs b/edb/native/pgrust/src/protocol/gen.rs similarity index 100% rename from edb/server/pgrust/src/protocol/gen.rs rename to edb/native/pgrust/src/protocol/gen.rs diff --git a/edb/server/pgrust/src/protocol/message_group.rs b/edb/native/pgrust/src/protocol/message_group.rs similarity index 100% rename from edb/server/pgrust/src/protocol/message_group.rs rename to edb/native/pgrust/src/protocol/message_group.rs diff --git a/edb/server/pgrust/src/protocol/mod.rs b/edb/native/pgrust/src/protocol/mod.rs similarity index 100% rename from edb/server/pgrust/src/protocol/mod.rs rename to edb/native/pgrust/src/protocol/mod.rs diff --git a/edb/server/pgrust/src/protocol/writer.rs b/edb/native/pgrust/src/protocol/writer.rs similarity index 100% rename from edb/server/pgrust/src/protocol/writer.rs rename to edb/native/pgrust/src/protocol/writer.rs diff --git a/edb/server/pgrust/src/python.rs b/edb/native/pgrust/src/python.rs similarity index 100% rename from edb/server/pgrust/src/python.rs rename to edb/native/pgrust/src/python.rs diff --git a/edb/server/pgrust/tests/edgedb_test_cases.rs b/edb/native/pgrust/tests/edgedb_test_cases.rs similarity index 100% rename from edb/server/pgrust/tests/edgedb_test_cases.rs rename to edb/native/pgrust/tests/edgedb_test_cases.rs diff --git a/edb/server/pgrust/tests/hardcore_host_tests_cases.rs b/edb/native/pgrust/tests/hardcore_host_tests_cases.rs similarity index 100% rename from edb/server/pgrust/tests/hardcore_host_tests_cases.rs rename to edb/native/pgrust/tests/hardcore_host_tests_cases.rs diff --git a/edb/server/pgrust/tests/libpq_test_cases.rs b/edb/native/pgrust/tests/libpq_test_cases.rs similarity index 100% rename from edb/server/pgrust/tests/libpq_test_cases.rs rename to edb/native/pgrust/tests/libpq_test_cases.rs diff --git a/edb/server/pgrust/tests/real_postgres.rs b/edb/native/pgrust/tests/real_postgres.rs similarity index 100% rename from edb/server/pgrust/tests/real_postgres.rs rename to edb/native/pgrust/tests/real_postgres.rs diff --git a/edb/server/pgrust/tests/test_util/dsn_libpq.rs b/edb/native/pgrust/tests/test_util/dsn_libpq.rs similarity index 100% rename from edb/server/pgrust/tests/test_util/dsn_libpq.rs rename to edb/native/pgrust/tests/test_util/dsn_libpq.rs diff --git a/edb/server/pgrust/tests/test_util/mod.rs b/edb/native/pgrust/tests/test_util/mod.rs similarity index 100% rename from edb/server/pgrust/tests/test_util/mod.rs rename to edb/native/pgrust/tests/test_util/mod.rs diff --git a/edb/native/pyo3_util/Cargo.toml b/edb/native/pyo3_util/Cargo.toml new file mode 100644 index 000000000000..f0930c0280e8 --- /dev/null +++ b/edb/native/pyo3_util/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "pyo3_util" +version = "0.1.0" +license = "MIT/Apache-2.0" +authors = ["MagicStack Inc. "] +edition = "2021" + +[lint] +workspace = true + +[dependencies] +pyo3 = { workspace = true } +tokio.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true +nix = { version = "0.29", features = ["fs"] } + +scopeguard = "1" +thiserror = "1" + +[lib] +crate-type = ["lib", "cdylib"] +path = "src/lib.rs" diff --git a/edb/native/pyo3_util/src/lib.rs b/edb/native/pyo3_util/src/lib.rs new file mode 100644 index 000000000000..31348d2fdacb --- /dev/null +++ b/edb/native/pyo3_util/src/lib.rs @@ -0,0 +1 @@ +pub mod logging; diff --git a/edb/native/pyo3_util/src/logging.rs b/edb/native/pyo3_util/src/logging.rs new file mode 100644 index 000000000000..de887017688b --- /dev/null +++ b/edb/native/pyo3_util/src/logging.rs @@ -0,0 +1,335 @@ +use pyo3::prelude::*; +use pyo3::{types::PyAnyMethods, Py, PyAny, PyResult, Python}; +use scopeguard::defer; +use std::cell::RefCell; +use std::collections::HashMap; +use std::sync::Mutex; +use std::sync::OnceLock; +use tracing::subscriber::DefaultGuard; +use tracing::Dispatch; +use tracing_subscriber::filter::LevelFilter; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::prelude::*; + +static EDGEDB_RUST_PYTHON_LOGGER_DEBUG: OnceLock = OnceLock::new(); +static LOGGER: OnceLock> = OnceLock::new(); + +fn is_debug_enabled() -> bool { + *EDGEDB_RUST_PYTHON_LOGGER_DEBUG.get_or_init(|| { + std::env::var("EDGEDB_RUST_PYTHON_LOGGER_DEBUG") + .map(|v| v == "1") + .unwrap_or(false) + }) +} + +/// A useful tool for debugging logging. +#[macro_export] +macro_rules! debug_log_method { + ($method_name:expr, $($arg:tt)*) => { + if is_debug_enabled() { + debug_log!($($arg)*); + defer! { + debug_log!("{} exited", $method_name); + } + } + }; +} + +/// A simple debug logging macro that prints to stderr if debug logging is enabled. +#[macro_export] +macro_rules! debug_log { + ($($arg:tt)*) => { + if is_debug_enabled() { + eprint!("LOGGING [{}]: ", std::process::id()); + eprintln!($($arg)*); + } + }; +} + +/// Provides a logger for the given Rust package and Python logger. Initializes +/// the logger if it is not already initialized. +/// +/// This is the public interface for the Python/Rust logging integration system. +pub fn get_logger(py: Python, rust_package: &str, python_logger: &str) -> PyResult { + debug_log_method!( + "get_logger", + "get_logger called with rust_package: {}, python_logger: {}", + rust_package, + python_logger + ); + LoggingGuard::new(rust_package.to_string(), python_logger.to_string(), py) +} + +/// Initializes logging for the current thread. This function should be called +/// at the start of any new thread that needs to use logging. +/// +/// Important: logging from threads requires taking the GIL and may have +/// performance impacts. +pub fn initialize_logging_in_thread() { + debug_log_method!( + "initialize_logging_in_thread", + "Initializing logging in thread" + ); + with_logger(|logger_bridge| { + thread_local! { + static GUARD: RefCell> = RefCell::new(None); + } + GUARD.with(|g| { + let dispatch = &logger_bridge + .dispatch + .as_ref() + .expect("LoggerBridge dispatch is not initialized") + .clone(); + *g.borrow_mut() = Some(tracing::dispatcher::set_default(dispatch)); + }); + }); +} + +fn with_logger(f: F) -> R +where + F: FnOnce(&mut LoggerBridge) -> R, +{ + LOGGER.get_or_init(|| Mutex::new(LoggerBridge::default())); + let mut guard = LOGGER.get().unwrap().lock().unwrap(); + f(&mut guard) +} + +fn python_to_rust_level(level: i32) -> LevelFilter { + match level { + ..10 => LevelFilter::TRACE, + 10 => LevelFilter::DEBUG, + 11..=20 => LevelFilter::INFO, + 21..=30 => LevelFilter::WARN, + 31..=40 => LevelFilter::ERROR, + _ => LevelFilter::OFF, + } +} + +fn log(py: Python, logger: &Py, event: &tracing::Event) { + debug_log_method!("log", "log function called"); + let mut message = format!("[{}] ", event.metadata().target()); + #[derive(Default)] + struct Visitor(String); + impl tracing::field::Visit for Visitor { + fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { + if field.name() == "message" { + self.0 += &format!("{value:?} "); + } else { + self.0 += &format!("{}={:?} ", field.name(), value) + } + } + } + + let mut visitor = Visitor::default(); + event.record(&mut visitor); + message += &visitor.0; + + let log_method = match *event.metadata().level() { + tracing::Level::TRACE | tracing::Level::DEBUG => "debug", + tracing::Level::INFO => "info", + tracing::Level::WARN => "warning", + tracing::Level::ERROR => "error", + }; + + if let Ok(log) = logger.getattr(py, log_method) { + let _ = log.call1(py, (message,)); + } +} + +struct LoggerInfo { + python_logger: Py, + level: LevelFilter, +} + +impl LoggerInfo { + fn clone_ref(&self, py: Python<'_>) -> Self { + LoggerInfo { + python_logger: self.python_logger.clone_ref(py), + level: self.level, + } + } +} + +#[derive(Default)] +struct LoggerBridge { + loggers: HashMap, + dispatch: Option, + guard: Option, +} + +impl tracing_subscriber::Layer for LoggerBridge { + fn on_event( + &self, + event: &tracing::Event<'_>, + _ctx: tracing_subscriber::layer::Context<'_, S>, + ) { + debug_log_method!("on_event", "LoggerBridge on_event called"); + Python::with_gil(|py| { + for (_, logger_info) in &self.loggers { + log(py, &logger_info.python_logger, event); + } + }); + } +} + +impl tracing_subscriber::layer::Filter for LoggerBridge { + fn enabled( + &self, + metadata: &tracing::Metadata<'_>, + _: &tracing_subscriber::layer::Context<'_, S>, + ) -> bool { + debug_log_method!( + "enabled", + "LoggerBridge enabled check for target: {} and level {:?}", + metadata.target(), + metadata.level() + ); + let crate_name = metadata + .target() + .split("::") + .next() + .unwrap_or(metadata.target()); + let result = if let Some(info) = self.loggers.get(crate_name) { + debug_log!("Testing against level: {:?}", info.level); + metadata.level() <= &info.level + } else { + debug_log!("No logger found for crate: {}", crate_name); + false + }; + result + } +} + +impl LoggerBridge { + fn clone_ref(&self, py: Python<'_>) -> Self { + LoggerBridge { + loggers: self + .loggers + .iter() + .map(|(k, v)| (k.clone(), v.clone_ref(py))) + .collect(), + dispatch: self.dispatch.clone(), + guard: None, + } + } + + fn add_logger( + &mut self, + rust_package: String, + python_logger: String, + py: Python, + ) -> PyResult<()> { + debug_log_method!( + "add_logger", + "LoggerBridge add_logger called for rust_package: {}, python_logger: {}", + rust_package, + python_logger + ); + let logging = py.import_bound("logging")?; + let logger = logging.getattr("getLogger")?.call1((python_logger,))?; + let level = logger + .getattr("getEffectiveLevel")? + .call0()? + .extract::()?; + let level_filter = python_to_rust_level(level); + debug_log!( + "Logger level updated - Python level: {}, Rust level: {:?}", + level, + level_filter + ); + self.loggers.insert( + rust_package, + LoggerInfo { + python_logger: logger.into(), + level: level_filter, + }, + ); + Ok(()) + } + + fn remove_logger(&mut self, rust_package: &str) { + debug_log_method!( + "remove_logger", + "LoggerBridge remove_logger called for rust_package: {}", + rust_package + ); + self.loggers.remove(rust_package); + } + + fn update_subscriber(&mut self, py: Python) -> PyResult<()> { + debug_log_method!("update_subscriber", "LoggerBridge update_subscriber called"); + let subscriber = + tracing_subscriber::registry().with(self.clone_ref(py).with_filter(self.clone_ref(py))); + let dispatch = Dispatch::new(subscriber); + self.dispatch = Some(dispatch.clone()); + self.guard.take(); + self.guard = Some(tracing::dispatcher::set_default(&dispatch)); + Ok(()) + } + + fn refresh(&mut self, py: Python) -> PyResult<()> { + debug_log_method!("refresh", "LoggerBridge refresh called"); + for (_, logger_info) in &mut self.loggers { + let logger = logger_info.python_logger.clone_ref(py); + let level = logger + .getattr(py, "getEffectiveLevel")? + .call0(py)? + .extract::(py)?; + let level_filter = python_to_rust_level(level); + debug_log!( + "Logger level updated - Python level: {}, Rust level: {:?}", + level, + level_filter + ); + logger_info.level = level_filter; + } + self.update_subscriber(py) + } +} + +#[pyclass] +pub struct LoggingGuard { + rust_package: String, +} + +impl Drop for LoggingGuard { + fn drop(&mut self) { + debug_log_method!( + "drop", + "LoggingGuard drop called for rust_package: {}", + self.rust_package + ); + let _ = Python::with_gil(|py| { + with_logger(|logger| { + logger.remove_logger(&self.rust_package); + logger.update_subscriber(py) + }) + }); + } +} + +impl LoggingGuard { + fn new(rust_package: String, python_logger: String, py: Python) -> PyResult { + debug_log_method!( + "new", + "LoggingGuard new called for rust_package: {}, python_logger: {}", + rust_package, + python_logger + ); + let result = with_logger(|logger| { + logger.add_logger(rust_package.clone(), python_logger, py)?; + logger.update_subscriber(py)?; + Ok(LoggingGuard { rust_package }) + }); + result + } +} + +#[pymethods] +impl LoggingGuard { + pub fn refresh(&self) -> PyResult<()> { + debug_log_method!("refresh", "LoggingGuard refresh called"); + let result = Python::with_gil(|py| with_logger(|logger| logger.refresh(py))); + result + } +} diff --git a/edb/server/edbrust/edbrust-util/Cargo.toml b/edb/server/_rust_native/Cargo.toml similarity index 59% rename from edb/server/edbrust/edbrust-util/Cargo.toml rename to edb/server/_rust_native/Cargo.toml index 6c6866a603d8..e4bd8d0e2185 100644 --- a/edb/server/edbrust/edbrust-util/Cargo.toml +++ b/edb/server/_rust_native/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "edbrust-rust" +name = "edbrust" version = "0.1.0" license = "MIT/Apache-2.0" authors = ["MagicStack Inc. "] @@ -14,6 +14,10 @@ python_extension = ["pyo3/extension-module", "pyo3/serde"] [dependencies] pyo3 = { workspace = true, optional = true } tokio.workspace = true +pyo3_util.workspace = true +conn_pool = { workspace = true, features = [ "python_extension" ] } +pgrust = { workspace = true, features = [ "python_extension" ] } +http = { workspace = true, features = [ "python_extension" ] } [lib] crate-type = ["lib", "cdylib"] diff --git a/edb/server/_rust_native/src/lib.rs b/edb/server/_rust_native/src/lib.rs new file mode 100644 index 000000000000..bae2b70cdfab --- /dev/null +++ b/edb/server/_rust_native/src/lib.rs @@ -0,0 +1,31 @@ +use pyo3::{ + pyfunction, pymodule, + types::{PyAnyMethods, PyModule, PyModuleMethods}, + wrap_pyfunction, Bound, IntoPy, Py, PyAny, PyResult, Python, +}; + +#[pymodule(name = "module")] +fn _rust_native(py: Python, m: &Bound) -> PyResult<()> { + let child_module = PyModule::new_bound(py, "edb.server._rust_native.module._conn_pool")?; + conn_pool::python::_conn_pool(py, &child_module)?; + m.add("_conn_pool", &child_module)?; + py.import_bound("sys")? + .getattr("modules")? + .set_item("edb.server._rust_native.module._conn_pool", child_module)?; + + let child_module = PyModule::new_bound(py, "edb.server._rust_native.module._pg_rust")?; + pgrust::python::_pg_rust(py, &child_module)?; + m.add("_pg_rust", &child_module)?; + py.import_bound("sys")? + .getattr("modules")? + .set_item("edb.server._rust_native.module._pg_rust", child_module)?; + + let child_module = PyModule::new_bound(py, "edb.server._rust_native.module._http")?; + http::python::_http(py, &child_module)?; + m.add("_http", &child_module)?; + py.import_bound("sys")? + .getattr("modules")? + .set_item("edb.server._rust_native.module._http", child_module)?; + + Ok(()) +} diff --git a/edb/server/edbrust/Cargo.toml b/edb/server/edbrust/Cargo.toml deleted file mode 100644 index c5bbd5a77ce9..000000000000 --- a/edb/server/edbrust/Cargo.toml +++ /dev/null @@ -1,40 +0,0 @@ -[package] -name = "edbrust" -version = "0.1.0" -license = "MIT/Apache-2.0" -authors = ["MagicStack Inc. "] -edition = "2021" - -[lint] -workspace = true - -[features] -python_extension = ["pyo3/extension-module", "pyo3/serde"] - -[dependencies] -pyo3 = { workspace = true, optional = true } -tokio.workspace = true - -[dependencies.derive_more] -version = "1.0.0-beta.6" -features = ["full"] - -[dev-dependencies] -pretty_assertions = "1.2.0" -test-log = { version = "0", features = ["trace"] } -anyhow = "1" -rstest = "0" -statrs = "0" -lru = "0" -byteorder = "1.5" -clap = "4" -clap_derive = "4" -hex-literal = "0.4" - -[dev-dependencies.tokio] -version = "1" -features = ["test-util"] - -[lib] -crate-type = ["lib", "cdylib"] -path = "src/lib.rs" diff --git a/edb/server/edbrust/edbrust-util/src/lib.rs b/edb/server/edbrust/edbrust-util/src/lib.rs deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/edb/server/edbrust/src/lib.rs b/edb/server/edbrust/src/lib.rs deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/setup.py b/setup.py index a33531125263..a69e044329a3 100644 --- a/setup.py +++ b/setup.py @@ -1182,20 +1182,8 @@ def _version(): binding=setuptools_rust.Binding.PyO3, ), setuptools_rust.RustExtension( - "edb.server._conn_pool", - path="edb/server/conn_pool/Cargo.toml", - features=["python_extension"], - binding=setuptools_rust.Binding.PyO3, - ), - setuptools_rust.RustExtension( - "edb.server._pg_rust", - path="edb/server/pgrust/Cargo.toml", - features=["python_extension"], - binding=setuptools_rust.Binding.PyO3, - ), - setuptools_rust.RustExtension( - "edb.server._http", - path="edb/server/http/Cargo.toml", + "edb.server._rust_native", + path="edb/server/_rust_native/Cargo.toml", features=["python_extension"], binding=setuptools_rust.Binding.PyO3, ),