Skip to content

Commit 7d36f96

Browse files
penbergpenberg
committed
Add timeout option to connect()
...including default busy timeout of 5 seconds like with Python SQLite module.
1 parent f621374 commit 7d36f96

File tree

2 files changed

+32
-16
lines changed

2 files changed

+32
-16
lines changed

‎src/lib.rs

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use pyo3::prelude::*;
55
use pyo3::types::{PyList, PyTuple};
66
use std::cell::{OnceCell, RefCell};
77
use std::sync::{Arc, OnceLock};
8+
use std::time::Duration;
89
use tokio::runtime::{Handle, Runtime};
910

1011
const LEGACY_TRANSACTION_CONTROL: i32 = -1;
@@ -37,10 +38,11 @@ fn is_remote_path(path: &str) -> bool {
3738

3839
#[pyfunction]
3940
#[cfg(not(Py_3_12))]
40-
#[pyo3(signature = (database, isolation_level="DEFERRED".to_string(), check_same_thread=true, uri=false, sync_url=None, sync_interval=None, auth_token="", encryption_key=None))]
41+
#[pyo3(signature = (database, timeout=5.0, isolation_level="DEFERRED".to_string(), check_same_thread=true, uri=false, sync_url=None, sync_interval=None, auth_token="", encryption_key=None))]
4142
fn connect(
4243
py: Python<'_>,
4344
database: String,
45+
timeout: f64,
4446
isolation_level: Option<String>,
4547
check_same_thread: bool,
4648
uri: bool,
@@ -52,6 +54,7 @@ fn connect(
5254
let conn = _connect_core(
5355
py,
5456
database,
57+
timeout,
5558
isolation_level,
5659
check_same_thread,
5760
uri,
@@ -65,10 +68,11 @@ fn connect(
6568

6669
#[pyfunction]
6770
#[cfg(Py_3_12)]
68-
#[pyo3(signature = (database, isolation_level="DEFERRED".to_string(), check_same_thread=true, uri=false, sync_url=None, sync_interval=None, auth_token="", encryption_key=None, autocommit = LEGACY_TRANSACTION_CONTROL))]
71+
#[pyo3(signature = (database, timeout=5.0, isolation_level="DEFERRED".to_string(), check_same_thread=true, uri=false, sync_url=None, sync_interval=None, auth_token="", encryption_key=None, autocommit = LEGACY_TRANSACTION_CONTROL))]
6972
fn connect(
7073
py: Python<'_>,
7174
database: String,
75+
timeout: f64,
7276
isolation_level: Option<String>,
7377
check_same_thread: bool,
7478
uri: bool,
@@ -81,6 +85,7 @@ fn connect(
8185
let mut conn = _connect_core(
8286
py,
8387
database,
88+
timeout,
8489
isolation_level.clone(),
8590
check_same_thread,
8691
uri,
@@ -104,6 +109,7 @@ fn connect(
104109
fn _connect_core(
105110
py: Python<'_>,
106111
database: String,
112+
timeout: f64,
107113
isolation_level: Option<String>,
108114
check_same_thread: bool,
109115
uri: bool,
@@ -130,8 +136,11 @@ fn _connect_core(
130136
match sync_url {
131137
Some(sync_url) => {
132138
let sync_interval = sync_interval.map(|i| std::time::Duration::from_secs_f64(i));
133-
let mut builder =
134-
libsql_core::Builder::new_remote_replica(database, sync_url, auth_token.to_string());
139+
let mut builder = libsql_core::Builder::new_remote_replica(
140+
database,
141+
sync_url,
142+
auth_token.to_string(),
143+
);
135144
if let Some(encryption_config) = encryption_config {
136145
builder = builder.encryption_config(encryption_config);
137146
}
@@ -158,6 +167,8 @@ fn _connect_core(
158167

159168
let autocommit = isolation_level.is_none() as i32;
160169
let conn = db.connect().map_err(to_py_err)?;
170+
let timeout = Duration::from_secs_f64(timeout);
171+
conn.busy_timeout(timeout).map_err(to_py_err)?;
161172
Ok(Connection {
162173
db,
163174
conn: RefCell::new(Some(Arc::new(ConnectionGuard {

‎tests/test_suite.py

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,12 @@
44
import sys
55
import libsql
66
import pytest
7+
import tempfile
78

9+
@pytest.mark.parametrize("provider", ["libsql", "sqlite"])
10+
def test_connection_timeout(provider):
11+
conn = connect(provider, ":memory:", timeout=1.0)
12+
conn.close()
813

914
@pytest.mark.parametrize("provider", ["libsql", "sqlite"])
1015
def test_connection_close(provider):
@@ -163,7 +168,7 @@ def test_commit_and_rollback(provider):
163168

164169
@pytest.mark.parametrize("provider", ["libsql", "sqlite"])
165170
def test_autocommit(provider):
166-
conn = connect(provider, ":memory:", None)
171+
conn = connect(provider, ":memory:", timeout=4, isolation_level=None)
167172
assert conn.isolation_level == None
168173
assert conn.in_transaction == False
169174
cur = conn.cursor()
@@ -182,7 +187,7 @@ def test_autocommit(provider):
182187
@pytest.mark.skipif(sys.version_info < (3, 12), reason="requires python3.12 or higher")
183188
def test_connection_autocommit(provider):
184189
# Test LEGACY_TRANSACTION_CONTROL (-1)
185-
conn = connect(provider, ":memory:", None, autocommit=-1)
190+
conn = connect(provider, ":memory:", timeout=5, isolation_level=None, autocommit=-1)
186191
assert conn.isolation_level is None
187192
assert conn.autocommit == -1
188193
cur = conn.cursor()
@@ -193,7 +198,7 @@ def test_connection_autocommit(provider):
193198
res = cur.execute("SELECT * FROM users")
194199
assert (1, "alice@example.com") == res.fetchone()
195200

196-
conn = connect(provider, ":memory:", isolation_level="DEFERRED", autocommit=-1)
201+
conn = connect(provider, ":memory:", timeout=5, isolation_level="DEFERRED", autocommit=-1)
197202
assert conn.isolation_level == "DEFERRED"
198203
assert conn.autocommit == -1
199204
cur = conn.cursor()
@@ -205,7 +210,7 @@ def test_connection_autocommit(provider):
205210
assert (1, "alice@example.com") == res.fetchone()
206211

207212
# Test autocommit Enabled (True)
208-
conn = connect(provider, ":memory:", None, autocommit=True)
213+
conn = connect(provider, ":memory:", timeout=5, isolation_level=None, autocommit=True)
209214
assert conn.isolation_level == None
210215
assert conn.autocommit == True
211216
cur = conn.cursor()
@@ -216,7 +221,7 @@ def test_connection_autocommit(provider):
216221
res = cur.execute("SELECT * FROM users")
217222
assert (1, "bob@example.com") == res.fetchone()
218223

219-
conn = connect(provider, ":memory:", isolation_level="DEFERRED", autocommit=True)
224+
conn = connect(provider, ":memory:", timeout=5, isolation_level="DEFERRED", autocommit=True)
220225
assert conn.isolation_level == "DEFERRED"
221226
assert conn.autocommit == True
222227
cur = conn.cursor()
@@ -228,7 +233,7 @@ def test_connection_autocommit(provider):
228233
assert (1, "bob@example.com") == res.fetchone()
229234

230235
# Test autocommit Disabled (False)
231-
conn = connect(provider, ":memory:", isolation_level="DEFERRED", autocommit=False)
236+
conn = connect(provider, ":memory:", timeout=5, isolation_level="DEFERRED", autocommit=False)
232237
assert conn.isolation_level == "DEFERRED"
233238
assert conn.autocommit == False
234239
cur = conn.cursor()
@@ -243,7 +248,7 @@ def test_connection_autocommit(provider):
243248

244249
# Test invalid autocommit value (should raise an error)
245250
with pytest.raises(ValueError):
246-
connect(provider, ":memory:", None, autocommit=999)
251+
connect(provider, ":memory:", timeout=5, isolation_level=None, autocommit=999)
247252

248253

249254
@pytest.mark.parametrize("provider", ["libsql", "sqlite"])
@@ -316,7 +321,7 @@ def test_int64(provider):
316321
assert [(1, 1099511627776)] == res.fetchall()
317322

318323

319-
def connect(provider, database, isolation_level="DEFERRED", autocommit=-1):
324+
def connect(provider, database, timeout=5, isolation_level="DEFERRED", autocommit=-1):
320325
if provider == "libsql-remote":
321326
from urllib import request
322327

@@ -332,21 +337,21 @@ def connect(provider, database, isolation_level="DEFERRED", autocommit=-1):
332337
if provider == "libsql":
333338
if sys.version_info < (3, 12):
334339
return libsql.connect(
335-
database, isolation_level=isolation_level
340+
database, timeout=timeout, isolation_level=isolation_level
336341
)
337342
else:
338343
if autocommit == -1:
339344
autocommit = libsql.LEGACY_TRANSACTION_CONTROL
340345
return libsql.connect(
341-
database, isolation_level=isolation_level, autocommit=autocommit
346+
database, timeout=timeout, isolation_level=isolation_level, autocommit=autocommit
342347
)
343348
if provider == "sqlite":
344349
if sys.version_info < (3, 12):
345-
return sqlite3.connect(database, isolation_level=isolation_level)
350+
return sqlite3.connect(database, timeout=timeout, isolation_level=isolation_level)
346351
else:
347352
if autocommit == -1:
348353
autocommit = sqlite3.LEGACY_TRANSACTION_CONTROL
349354
return sqlite3.connect(
350-
database, isolation_level=isolation_level, autocommit=autocommit
355+
database, timeout=timeout, isolation_level=isolation_level, autocommit=autocommit
351356
)
352357
raise Exception(f"Provider `{provider}` is not supported")

0 commit comments

Comments
 (0)