From cb69d8b0dd66f9e7e898ec1591bd10991947c688 Mon Sep 17 00:00:00 2001 From: JeanArhancet Date: Mon, 23 Dec 2024 17:05:21 +0100 Subject: [PATCH 1/3] feat(python): add in-memory mode --- bindings/python/build.rs | 1 - bindings/python/src/lib.rs | 86 ++++++++++++++----------- bindings/python/tests/hello.db | Bin 0 -> 12288 bytes bindings/python/tests/test_database.py | 14 ++++ 4 files changed, 64 insertions(+), 37 deletions(-) create mode 100644 bindings/python/tests/hello.db diff --git a/bindings/python/build.rs b/bindings/python/build.rs index 8f01c1a67..0475124bb 100644 --- a/bindings/python/build.rs +++ b/bindings/python/build.rs @@ -1,4 +1,3 @@ fn main() { pyo3_build_config::use_pyo3_cfgs(); - println!("cargo::rustc-check-cfg=cfg(allocator, values(\"default\", \"mimalloc\"))"); } diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 1b3514032..9d1cba7b3 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -4,6 +4,7 @@ use limbo_core::IO; use pyo3::prelude::*; use pyo3::types::PyList; use pyo3::types::PyTuple; +use std::cell::RefCell; use std::rc::Rc; use std::sync::{Arc, Mutex}; @@ -78,7 +79,7 @@ pub struct Cursor { #[pyo3(get)] rowcount: i64, - smt: Option>>, + smt: Option>>, } // SAFETY: The limbo_core crate guarantees that `Cursor` is thread-safe. @@ -90,26 +91,33 @@ impl Cursor { #[pyo3(signature = (sql, parameters=None))] pub fn execute(&mut self, sql: &str, parameters: Option>) -> Result { let stmt_is_dml = stmt_is_dml(sql); + let stmt_is_ddl = stmt_is_ddl(sql); - let conn_lock = - self.conn.conn.lock().map_err(|_| { - PyErr::new::("Failed to acquire connection lock") - })?; - - let statement = conn_lock.prepare(sql).map_err(|e| { + let statement = self.conn.conn.prepare(sql).map_err(|e| { PyErr::new::(format!("Failed to prepare statement: {:?}", e)) })?; - self.smt = Some(Arc::new(Mutex::new(statement))); + let stmt = Rc::new(RefCell::new(statement)); - // TODO: use stmt_is_dml to set rowcount - if stmt_is_dml { - return Err(PyErr::new::( - "DML statements (INSERT/UPDATE/DELETE) are not fully supported in this version", - ) - .into()); + // For DDL and DML statements, + // we need to execute the statement immediately + if stmt_is_ddl || stmt_is_dml { + loop { + match stmt.borrow_mut().step().map_err(|e| { + PyErr::new::(format!("Step error: {:?}", e)) + })? { + limbo_core::RowResult::IO => { + self.conn.io.run_once().map_err(|e| { + PyErr::new::(format!("IO error: {:?}", e)) + })?; + } + _ => break, + } + } } + self.smt = Some(stmt); + Ok(Cursor { smt: self.smt.clone(), conn: self.conn.clone(), @@ -121,11 +129,8 @@ impl Cursor { pub fn fetchone(&mut self, py: Python) -> Result> { if let Some(smt) = &self.smt { - let mut smt_lock = smt.lock().map_err(|_| { - PyErr::new::("Failed to acquire statement lock") - })?; loop { - match smt_lock.step().map_err(|e| { + match smt.borrow_mut().step().map_err(|e| { PyErr::new::(format!("Step error: {:?}", e)) })? { limbo_core::StepResult::Row(row) => { @@ -157,14 +162,9 @@ impl Cursor { pub fn fetchall(&mut self, py: Python) -> Result> { let mut results = Vec::new(); - if let Some(smt) = &self.smt { - let mut smt_lock = smt.lock().map_err(|_| { - PyErr::new::("Failed to acquire statement lock") - })?; - loop { - match smt_lock.step().map_err(|e| { + match smt.borrow_mut().step().map_err(|e| { PyErr::new::(format!("Step error: {:?}", e)) })? { limbo_core::StepResult::Row(row) => { @@ -221,11 +221,17 @@ fn stmt_is_dml(sql: &str) -> bool { sql.starts_with("INSERT") || sql.starts_with("UPDATE") || sql.starts_with("DELETE") } +fn stmt_is_ddl(sql: &str) -> bool { + let sql = sql.trim(); + let sql = sql.to_uppercase(); + sql.starts_with("CREATE") || sql.starts_with("ALTER") || sql.starts_with("DROP") +} + #[pyclass] #[derive(Clone)] pub struct Connection { - conn: Arc>>, - io: Arc, + conn: Rc, + io: Arc, } // SAFETY: The limbo_core crate guarantees that `Connection` is thread-safe. @@ -263,16 +269,24 @@ impl Connection { #[allow(clippy::arc_with_non_send_sync)] #[pyfunction] pub fn connect(path: &str) -> Result { - let io = Arc::new(limbo_core::PlatformIO::new().map_err(|e| { - PyErr::new::(format!("IO initialization failed: {:?}", e)) - })?); - let db = limbo_core::Database::open_file(io.clone(), path) - .map_err(|e| PyErr::new::(format!("Failed to open database: {:?}", e)))?; - let conn: Rc = db.connect(); - Ok(Connection { - conn: Arc::new(Mutex::new(conn)), - io, - }) + match path { + ":memory:" => { + let io: Arc = Arc::new(limbo_core::MemoryIO::new()?); + let db = limbo_core::Database::open_file(io.clone(), path).map_err(|e| { + PyErr::new::(format!("Failed to open database: {:?}", e)) + })?; + let conn: Rc = db.connect(); + Ok(Connection { conn, io }) + } + path => { + let io: Arc = Arc::new(limbo_core::PlatformIO::new()?); + let db = limbo_core::Database::open_file(io.clone(), path).map_err(|e| { + PyErr::new::(format!("Failed to open database: {:?}", e)) + })?; + let conn: Rc = db.connect(); + Ok(Connection { conn, io }) + } + } } fn row_to_py(py: Python, row: &limbo_core::Row) -> PyObject { diff --git a/bindings/python/tests/hello.db b/bindings/python/tests/hello.db new file mode 100644 index 0000000000000000000000000000000000000000..55538d0cffb2128facdeb01a7a07ec836d9c8ae1 GIT binary patch literal 12288 zcmeI#ze@u#6u|N1lp=+4-Lk&Xf(YU)cuI!iDSBpTr{km%4*COoSJA0|fPb6+kE54B z=kCY%L0)(X3FMRMy}smi>*am3T~)1ItFh8r-HIrsbT>}Aafnc#9*S-b^|8?E^7Chu z{;4G0t8}0K9y`E30tg_000IagfB*srAb-X|51Q0*~0R#|0 O009ILKmY**qQF0p!Z6AJ literal 0 HcmV?d00001 diff --git a/bindings/python/tests/test_database.py b/bindings/python/tests/test_database.py index 601c67135..a09eea08e 100644 --- a/bindings/python/tests/test_database.py +++ b/bindings/python/tests/test_database.py @@ -27,6 +27,20 @@ def test_fetchall_select_user_ids(provider): assert user_ids == [(1,), (2,)] +@pytest.mark.parametrize("provider", ["sqlite3", "limbo"]) +def test_in_memory_fetchone_select_all_users(provider): + conn = connect(provider, ":memory:") + cursor = conn.cursor() + cursor.execute("CREATE TABLE users (id INT PRIMARY KEY, username TEXT)") + cursor.execute("INSERT INTO users VALUES (1, 'alice')") + + cursor.execute("SELECT * FROM users") + + alice = cursor.fetchone() + assert alice + assert alice == (1, "alice") + + @pytest.mark.parametrize("provider", ["sqlite3", "limbo"]) def test_fetchone_select_all_users(provider): conn = connect(provider, "tests/database.db") From 2a0402ce7f0510b645bcad99c09120028d1977aa Mon Sep 17 00:00:00 2001 From: JeanArhancet Date: Sat, 28 Dec 2024 17:17:57 +0100 Subject: [PATCH 2/3] fix: python lint --- bindings/python/src/lib.rs | 2 +- bindings/python/tests/hello.db | Bin 12288 -> 0 bytes bindings/python/tests/test_database.py | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) delete mode 100644 bindings/python/tests/hello.db diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 9d1cba7b3..7c2658a9a 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -106,7 +106,7 @@ impl Cursor { match stmt.borrow_mut().step().map_err(|e| { PyErr::new::(format!("Step error: {:?}", e)) })? { - limbo_core::RowResult::IO => { + limbo_core::StepResult::IO => { self.conn.io.run_once().map_err(|e| { PyErr::new::(format!("IO error: {:?}", e)) })?; diff --git a/bindings/python/tests/hello.db b/bindings/python/tests/hello.db deleted file mode 100644 index 55538d0cffb2128facdeb01a7a07ec836d9c8ae1..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 12288 zcmeI#ze@u#6u|N1lp=+4-Lk&Xf(YU)cuI!iDSBpTr{km%4*COoSJA0|fPb6+kE54B z=kCY%L0)(X3FMRMy}smi>*am3T~)1ItFh8r-HIrsbT>}Aafnc#9*S-b^|8?E^7Chu z{;4G0t8}0K9y`E30tg_000IagfB*srAb-X|51Q0*~0R#|0 O009ILKmY**qQF0p!Z6AJ diff --git a/bindings/python/tests/test_database.py b/bindings/python/tests/test_database.py index a09eea08e..63241f19a 100644 --- a/bindings/python/tests/test_database.py +++ b/bindings/python/tests/test_database.py @@ -33,7 +33,7 @@ def test_in_memory_fetchone_select_all_users(provider): cursor = conn.cursor() cursor.execute("CREATE TABLE users (id INT PRIMARY KEY, username TEXT)") cursor.execute("INSERT INTO users VALUES (1, 'alice')") - + cursor.execute("SELECT * FROM users") alice = cursor.fetchone() From 9a70dc8f78372de143a5919af3d3b21090be19dc Mon Sep 17 00:00:00 2001 From: JeanArhancet Date: Mon, 30 Dec 2024 10:22:36 +0100 Subject: [PATCH 3/3] fix: clippy error --- bindings/python/src/lib.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 7c2658a9a..764fcfcf5 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -1,12 +1,11 @@ use anyhow::Result; use errors::*; -use limbo_core::IO; use pyo3::prelude::*; use pyo3::types::PyList; use pyo3::types::PyTuple; use std::cell::RefCell; use std::rc::Rc; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; mod errors;