Merge pull request #37 from tursodatabase/lucio/add-nix-support

Fix no reactor running/add nix support/add ctrl+c for long running fn

Author: penberg

README.md

@@ -101,6 +101,12 @@ source .env/bin/activate
 pip3 install maturin pyperf pytest
 ```
 
+Or you can use NIX to drop you into a shell with everything installed
+
+```
+nix-shell
+```
+
 Build the development version and use it:
 
 ```

shell.nix

@@ -0,0 +1,17 @@
+{ pkgs ? import <nixpkgs> { } }:
+(pkgs.buildFHSUserEnv {
+  name = "pipzone";
+  targetPkgs = pkgs: (with pkgs; [
+    python39
+    python39Packages.pip
+    python39Packages.virtualenv
+    python39Packages.pytest
+    python39Packages.pyperf
+    maturin
+  ]);
+  runScript = "bash";
+  profile = ''
+    virtualenv .venv
+    source .venv/bin/activate
+  '';
+}).env

src/lib.rs

@@ -21,6 +21,7 @@ fn is_remote_path(path: &str) -> bool {
 #[pyfunction]
 #[pyo3(signature = (database, isolation_level="DEFERRED".to_string(), check_same_thread=true, uri=false, sync_url=None, auth_token=""))]
 fn connect(
+    py: Python<'_>,
     database: String,
     isolation_level: Option<String>,
     check_same_thread: bool,
@@ -46,28 +47,58 @@ fn connect(
                     None,
                     None,
                 );
-                let result = rt.block_on(fut);
+                tokio::pin!(fut);
+                let result = rt.block_on(check_signals(py, fut));
                 result.map_err(to_py_err)?
             }
             None => libsql_core::Database::open(database).map_err(to_py_err)?,
         }
     };
     let autocommit = isolation_level.is_none();
     let conn = db.connect().map_err(to_py_err)?;
-    let conn = Arc::new(conn);
     Ok(Connection {
         db,
-        conn,
+        conn: Arc::new(ConnectionGuard {
+            conn: Some(conn),
+            handle: rt.handle().clone(),
+        }),
         rt,
         isolation_level,
         autocommit,
     })
 }
 
+// We need to add a drop guard that runs when we finally drop our
+// only reference to libsql_core::Connection. This is because when
+// hrana is enabled it needs access to the tokio api to spawn a close
+// call in the background. So this adds the ability that when drop is called
+// on ConnectionGuard it will drop the connection with a tokio context entered.
+struct ConnectionGuard {
+    conn: Option<libsql_core::Connection>,
+    handle: tokio::runtime::Handle,
+}
+
+impl std::ops::Deref for ConnectionGuard {
+    type Target = libsql_core::Connection;
+
+    fn deref(&self) -> &Self::Target {
+        &self.conn.as_ref().expect("Connection already dropped")
+    }
+}
+
+impl Drop for ConnectionGuard {
+    fn drop(&mut self) {
+        let _enter = self.handle.enter();
+        if let Some(conn) = self.conn.take() {
+            drop(conn);
+        }
+    }
+}
+
 #[pyclass]
 pub struct Connection {
     db: libsql_core::Database,
-    conn: Arc<libsql_core::Connection>,
+    conn: Arc<ConnectionGuard>,
     rt: tokio::runtime::Runtime,
     isolation_level: Option<String>,
     autocommit: bool,
@@ -91,8 +122,17 @@ impl Connection {
         })
     }
 
-    fn sync(self_: PyRef<'_, Self>) -> PyResult<()> {
-        self_.rt.block_on(self_.db.sync()).map_err(to_py_err)?;
+    fn sync(self_: PyRef<'_, Self>, py: Python<'_>) -> PyResult<()> {
+        let fut = {
+            let _enter = self_.rt.enter();
+            self_.db.sync()
+        };
+        tokio::pin!(fut);
+
+        self_
+            .rt
+            .block_on(check_signals(py, fut))
+            .map_err(to_py_err)?;
         Ok(())
     }
 
@@ -101,7 +141,7 @@ impl Connection {
         if !self_.conn.is_autocommit() {
             self_
                 .rt
-                .block_on(self_.conn.execute("COMMIT", ()))
+                .block_on(async { self_.conn.execute("COMMIT", ()).await })
                 .map_err(to_py_err)?;
         }
         Ok(())
@@ -112,7 +152,7 @@ impl Connection {
         if !self_.conn.is_autocommit() {
             self_
                 .rt
-                .block_on(self_.conn.execute("ROLLBACK", ()))
+                .block_on(async { self_.conn.execute("ROLLBACK", ()).await })
                 .map_err(to_py_err)?;
         }
         Ok(())
@@ -125,7 +165,7 @@ impl Connection {
     ) -> PyResult<Cursor> {
         let cursor = Connection::cursor(&self_)?;
         let rt = self_.rt.handle();
-        rt.block_on(execute(&cursor, sql, parameters))?;
+        rt.block_on(async { execute(&cursor, sql, parameters).await })?;
         Ok(cursor)
     }
 
@@ -139,7 +179,7 @@ impl Connection {
             let parameters = parameters.extract::<&PyTuple>()?;
             self_
                 .rt
-                .block_on(execute(&cursor, sql.clone(), Some(parameters)))?;
+                .block_on(async { execute(&cursor, sql.clone(), Some(parameters)).await })?;
         }
         Ok(cursor)
     }
@@ -160,7 +200,7 @@ pub struct Cursor {
     #[pyo3(get, set)]
     arraysize: usize,
     rt: tokio::runtime::Handle,
-    conn: Arc<libsql_core::Connection>,
+    conn: Arc<ConnectionGuard>,
     stmt: RefCell<Option<libsql_core::Statement>>,
     rows: RefCell<Option<libsql_core::Rows>>,
     rowcount: RefCell<i64>,
@@ -178,7 +218,9 @@ impl Cursor {
         sql: String,
         parameters: Option<&PyTuple>,
     ) -> PyResult<pyo3::PyRef<'a, Self>> {
-        self_.rt.block_on(execute(&self_, sql, parameters))?;
+        self_
+            .rt
+            .block_on(async { execute(&self_, sql, parameters).await })?;
         Ok(self_)
     }
 
@@ -191,7 +233,7 @@ impl Cursor {
             let parameters = parameters.extract::<&PyTuple>()?;
             self_
                 .rt
-                .block_on(execute(&self_, sql.clone(), Some(parameters)))?;
+                .block_on(async { execute(&self_, sql.clone(), Some(parameters)).await })?;
         }
         Ok(self_)
     }
@@ -246,19 +288,22 @@ impl Cursor {
                 // done before iterating.
                 if !*self_.done.borrow() {
                     for _ in 0..size {
-                        let row = self_.rt.block_on(rows.next()).map_err(to_py_err)?;
+                        let row = self_
+                            .rt
+                            .block_on(async { rows.next().await })
+                            .map_err(to_py_err)?;
                         match row {
                             Some(row) => {
                                 let row = convert_row(self_.py(), row, rows.column_count())?;
                                 elements.push(row.into());
                             }
                             None => {
                                 self_.done.replace(true);
-                                break
+                                break;
                             }
                         }
                     }
-                    }
+                }
                 Ok(Some(PyList::new(self_.py(), elements)))
             }
             None => Ok(None),
@@ -271,7 +316,10 @@ impl Cursor {
             Some(rows) => {
                 let mut elements: Vec<Py<PyAny>> = vec![];
                 loop {
-                    let row = self_.rt.block_on(rows.next()).map_err(to_py_err)?;
+                    let row = self_
+                        .rt
+                        .block_on(async { rows.next().await })
+                        .map_err(to_py_err)?;
                     match row {
                         Some(row) => {
                             let row = convert_row(self_.py(), row, rows.column_count())?;
@@ -393,3 +441,20 @@ fn libsql_experimental(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<Cursor>()?;
     Ok(())
 }
+
+async fn check_signals<F, R>(py: Python<'_>, mut fut: std::pin::Pin<&mut F>) -> R
+where
+    F: std::future::Future<Output = R>,
+{
+    loop {
+        tokio::select! {
+            out = &mut fut => {
+                break out;
+            }
+
+            _ = tokio::time::sleep(std::time::Duration::from_millis(300)) => {
+                py.check_signals().unwrap();
+            }
+        }
+    }
+}