feat: allow custom vfs during initialization

Author: darkskygit

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sqlite-vfs-http"
-version = "0.1.3"
+version = "0.1.4"
 edition = "2021"
 authors = ["DarkSky <darksky2048@gmail.com>"]
 license = "AGPL-3.0-only"

src/conn.rs

@@ -39,8 +39,8 @@ impl Connection {
         None
     }
 
-    fn init_with_url(url: &str) -> Result<(AtomicRuntime, usize), Error> {
-        let rt = AtomicRuntime::default();
+    fn init_with_url(url: &str, client: Option<Client>) -> Result<(AtomicRuntime, usize), Error> {
+        let rt = AtomicRuntime::new(client);
         match Self::get_length(rt.clone(), url.to_string()) {
             Some(size) if size != 0 => Ok((rt, size)),
             _ => {
@@ -53,8 +53,13 @@ impl Connection {
         }
     }
 
-    pub fn new(url: &str, block_size: usize, download_threshold: usize) -> Result<Self, Error> {
-        let (rt, length) = Self::init_with_url(url)?;
+    pub fn new(
+        url: &str,
+        client: Option<Client>,
+        block_size: usize,
+        download_threshold: usize,
+    ) -> Result<Self, Error> {
+        let (rt, length) = Self::init_with_url(url, client)?;
         let buffer = LazyBuffer::new(
             length,
             block_size,

src/lib.rs

@@ -5,7 +5,8 @@ mod vfs;
 
 use buffer::LazyBuffer;
 use conn::Connection;
-use sqlite_vfs::register;
+use reqwest::Client;
+use sqlite_vfs::{register, RegisterError};
 use std::sync::{Arc, Once, RwLock};
 use utils::AtomicRuntime;
 use vfs::HttpVfs;
@@ -16,6 +17,8 @@ pub const HTTP_VFS: &str = "http";
 pub struct HttpVfsRegister {
     /// how many pages in block, default is 8MB, 2048 pages
     block_size: usize,
+    /// default client
+    client: Option<Client>,
     /// read the first few pages of each block without downloading the entire block
     download_threshold: usize,
     /// sqlite's page size is 4KB by default
@@ -25,6 +28,7 @@ pub struct HttpVfsRegister {
 impl HttpVfsRegister {
     pub fn new() -> Self {
         Self {
+            client: None,
             block_size: SQLITE_PAGE_SIZE * 1024 * 2,
             download_threshold: 0,
             page_size: SQLITE_PAGE_SIZE,
@@ -38,6 +42,13 @@ impl HttpVfsRegister {
         }
     }
 
+    pub fn with_client(self, client: Client) -> Self {
+        Self {
+            client: Some(client),
+            ..self
+        }
+    }
+
     /// Set how many page read don't download full block
     pub fn with_download_threshold(self, page_num: usize) -> Self {
         Self {
@@ -50,21 +61,33 @@ impl HttpVfsRegister {
         Self { page_size, ..self }
     }
 
-    pub fn register(self) {
-        const ONCE: Once = Once::new();
-
+    pub fn register(self) -> Result<(), RegisterError> {
         let vfs_instance = HttpVfs {
+            client: self.client,
             block_size: self.block_size,
             download_threshold: self.download_threshold,
         };
-
-        ONCE.call_once(|| {
-            let _ = register(HTTP_VFS, vfs_instance, true);
-        })
+        register(HTTP_VFS, vfs_instance, true)
     }
 }
 
+/// register http vfs, use `Once` internally to ensure only register once
 #[inline(always)]
 pub fn register_http_vfs() {
-    HttpVfsRegister::new().register();
+    const ONCE: Once = Once::new();
+
+    ONCE.call_once(|| {
+        let _ = HttpVfsRegister::new().register();
+    })
+}
+
+/// register http vfs with custom client
+/// use `Once` internally to ensure only register once
+#[inline(always)]
+pub fn register_http_vfs_with_custom(cb: impl FnOnce(HttpVfsRegister) -> HttpVfsRegister) {
+    const ONCE: Once = Once::new();
+
+    ONCE.call_once(|| {
+        let _ = cb(HttpVfsRegister::new()).register();
+    })
 }

src/utils.rs

@@ -9,10 +9,10 @@ pub struct AtomicRuntime {
     rt: Arc<RwLock<Option<Runtime>>>,
 }
 
-impl Default for AtomicRuntime {
-    fn default() -> Self {
+impl AtomicRuntime {
+    pub fn new(client: Option<Client>) -> Self {
         Self {
-            client: Client::new(),
+            client: client.unwrap_or_else(|| Client::new()),
             rt: Arc::new(RwLock::new(Runtime::new().ok())),
         }
     }

src/vfs.rs

@@ -7,6 +7,7 @@ use std::{
 };
 
 pub struct HttpVfs {
+    pub(crate) client: Option<Client>,
     pub(crate) block_size: usize,
     pub(crate) download_threshold: usize,
 }
@@ -24,6 +25,7 @@ impl Vfs for HttpVfs {
 
         Ok(Connection::new(
             db,
+            self.client.clone(),
             self.block_size,
             self.download_threshold,
         )?)