From ae2e5db4b2532516f34886963ada75f60264ae29 Mon Sep 17 00:00:00 2001
From: Julius de Bruijn <bruijn@prisma.io>
Date: Tue, 4 Aug 2020 18:31:56 +0200
Subject: [PATCH] SQLite: Allow setting `busy_timeout` in options

---
 sqlx-core/src/sqlite/connection/establish.rs | 32 ++++++++++++--------
 sqlx-core/src/sqlite/options/mod.rs          | 13 +++++++-
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/sqlx-core/src/sqlite/connection/establish.rs b/sqlx-core/src/sqlite/connection/establish.rs
index 6fd765d6..954cac4f 100644
--- a/sqlx-core/src/sqlite/connection/establish.rs
+++ b/sqlx-core/src/sqlite/connection/establish.rs
@@ -1,13 +1,3 @@
-use std::io;
-use std::ptr::{null, null_mut};
-
-use libsqlite3_sys::{
-    sqlite3_busy_timeout, sqlite3_extended_result_codes, sqlite3_open_v2, SQLITE_OK,
-    SQLITE_OPEN_CREATE, SQLITE_OPEN_MEMORY, SQLITE_OPEN_NOMUTEX, SQLITE_OPEN_PRIVATECACHE,
-    SQLITE_OPEN_READONLY, SQLITE_OPEN_READWRITE,
-};
-use sqlx_rt::blocking;
-
 use crate::error::Error;
 use crate::sqlite::connection::handle::ConnectionHandle;
 use crate::sqlite::statement::StatementWorker;
@@ -15,6 +5,17 @@ use crate::{
     common::StatementCache,
     sqlite::{SqliteConnectOptions, SqliteConnection, SqliteError},
 };
+use libsqlite3_sys::{
+    sqlite3_busy_timeout, sqlite3_extended_result_codes, sqlite3_open_v2, SQLITE_OK,
+    SQLITE_OPEN_CREATE, SQLITE_OPEN_MEMORY, SQLITE_OPEN_NOMUTEX, SQLITE_OPEN_PRIVATECACHE,
+    SQLITE_OPEN_READONLY, SQLITE_OPEN_READWRITE,
+};
+use sqlx_rt::blocking;
+use std::io;
+use std::{
+    convert::TryFrom,
+    ptr::{null, null_mut},
+};
 
 pub(crate) async fn establish(options: &SqliteConnectOptions) -> Result<SqliteConnection, Error> {
     let mut filename = options
@@ -48,6 +49,8 @@ pub(crate) async fn establish(options: &SqliteConnectOptions) -> Result<SqliteCo
         flags |= SQLITE_OPEN_MEMORY;
     }
 
+    let busy_timeout = options.busy_timeout;
+
     let handle = blocking!({
         let mut handle = null_mut();
 
@@ -85,8 +88,13 @@ pub(crate) async fn establish(options: &SqliteConnectOptions) -> Result<SqliteCo
         // This causes SQLite to automatically sleep in increasing intervals until the time
         // when there is something locked during [sqlite3_step]. This is sync. but we only
         // run [sqlite3_step] in [blocking!] so its okay.
-        // TODO: Allow this timeout to be configured in SqliteOptions
-        status = unsafe { sqlite3_busy_timeout(handle.0.as_ptr(), 5000) };
+        //
+        // We also need to convert the u128 value to i32, checking we're not overflowing.
+        let ms =
+            i32::try_from(busy_timeout.as_millis()).expect("Given busy timeout value is too big.");
+
+        status = unsafe { sqlite3_busy_timeout(handle.0.as_ptr(), ms) };
+
         if status != SQLITE_OK {
             return Err(Error::Database(Box::new(SqliteError::new(handle.as_ptr()))));
         }
diff --git a/sqlx-core/src/sqlite/options/mod.rs b/sqlx-core/src/sqlite/options/mod.rs
index fb2548ca..fce09967 100644
--- a/sqlx-core/src/sqlite/options/mod.rs
+++ b/sqlx-core/src/sqlite/options/mod.rs
@@ -5,7 +5,7 @@ mod journal_mode;
 mod parse;
 
 pub use journal_mode::SqliteJournalMode;
-use std::borrow::Cow;
+use std::{borrow::Cow, time::Duration};
 
 /// Options and flags which can be used to configure a SQLite connection.
 ///
@@ -49,6 +49,7 @@ pub struct SqliteConnectOptions {
     pub(crate) journal_mode: SqliteJournalMode,
     pub(crate) foreign_keys: bool,
     pub(crate) statement_cache_capacity: usize,
+    pub(crate) busy_timeout: Duration,
 }
 
 impl Default for SqliteConnectOptions {
@@ -67,6 +68,7 @@ impl SqliteConnectOptions {
             foreign_keys: true,
             statement_cache_capacity: 100,
             journal_mode: SqliteJournalMode::Wal,
+            busy_timeout: Duration::from_secs(5),
         }
     }
 
@@ -119,4 +121,13 @@ impl SqliteConnectOptions {
         self.statement_cache_capacity = capacity;
         self
     }
+
+    /// Sets a timeout value to wait when the database is locked, before
+    /// returning a busy timeout error.
+    ///
+    /// The default busy timeout is 5 seconds.
+    pub fn busy_timeout(mut self, timeout: Duration) -> Self {
+        self.busy_timeout = timeout;
+        self
+    }
 }