|
6 | 6 | //! The C shim provides symbol shims for `sqlite3_vfs_register` etc. that route |
7 | 7 | //! through the extension API table, so `sqlite_vfs::register()` works correctly |
8 | 8 | //! inside a loadable extension. |
| 9 | +//! |
| 10 | +//! ## VFS selection |
| 11 | +//! |
| 12 | +//! If `TURBOLITE_BUCKET` is set, registers a **tiered S3 VFS** (requires the |
| 13 | +//! `tiered` feature). Otherwise, registers a **local compressed VFS**. |
| 14 | +//! |
| 15 | +//! ### Environment variables (tiered mode) |
| 16 | +//! |
| 17 | +//! | Variable | Required | Default | Description | |
| 18 | +//! |---|---|---|---| |
| 19 | +//! | `TURBOLITE_BUCKET` | yes | — | S3 bucket name | |
| 20 | +//! | `TURBOLITE_PREFIX` | no | `"turbolite"` | S3 key prefix | |
| 21 | +//! | `TURBOLITE_CACHE_DIR` | no | `"/tmp/turbolite"` | Local cache directory | |
| 22 | +//! | `TURBOLITE_ENDPOINT_URL` | no | — | Custom S3 endpoint (Tigris, MinIO) | |
| 23 | +//! | `TURBOLITE_REGION` | no | — | AWS region | |
| 24 | +//! | `TURBOLITE_PREFETCH_THREADS` | no | `num_cpus + 1` | Prefetch worker threads | |
| 25 | +//! | `TURBOLITE_COMPRESSION_LEVEL` | no | `3` | Zstd level 1-22 | |
| 26 | +//! | `TURBOLITE_READ_ONLY` | no | `false` | Open in read-only mode | |
| 27 | +//! |
| 28 | +//! Falls back to `AWS_ENDPOINT_URL` / `AWS_REGION` if the `TURBOLITE_` variants |
| 29 | +//! are not set. |
9 | 30 |
|
10 | | -use std::path::PathBuf; |
11 | 31 | use std::sync::atomic::{AtomicBool, Ordering}; |
12 | 32 |
|
13 | | -use crate::CompressedVfs; |
14 | | - |
15 | | -/// Track whether the VFS has already been registered (idempotent load). |
16 | 33 | static VFS_REGISTERED: AtomicBool = AtomicBool::new(false); |
17 | 34 |
|
18 | 35 | /// Called from C entry point (`sqlite3_turbolite_init` in ext_entry.c). |
19 | | -/// |
20 | | -/// Registers a compressed VFS named "turbolite" with zstd level 3. |
21 | | -/// The base directory is "." — SQLite passes full paths to the VFS, so |
22 | | -/// relative base dir means paths are used as-given by the host application. |
23 | | -/// |
24 | 36 | /// Returns 0 on success, 1 on error. Idempotent: second call is a no-op. |
25 | 37 | #[no_mangle] |
26 | 38 | pub extern "C" fn turbolite_ext_register_vfs() -> std::os::raw::c_int { |
27 | 39 | if VFS_REGISTERED.swap(true, Ordering::SeqCst) { |
28 | | - // Already registered — idempotent success. |
29 | 40 | return 0; |
30 | 41 | } |
31 | 42 |
|
32 | | - let vfs = CompressedVfs::new(PathBuf::from("."), 3); |
33 | | - match crate::register("turbolite", vfs) { |
| 43 | + let result = if std::env::var("TURBOLITE_BUCKET").is_ok() { |
| 44 | + register_tiered() |
| 45 | + } else { |
| 46 | + register_local() |
| 47 | + }; |
| 48 | + |
| 49 | + match result { |
34 | 50 | Ok(()) => 0, |
35 | 51 | Err(_) => { |
36 | 52 | VFS_REGISTERED.store(false, Ordering::SeqCst); |
37 | 53 | 1 |
38 | 54 | } |
39 | 55 | } |
40 | 56 | } |
| 57 | + |
| 58 | +fn register_local() -> Result<(), std::io::Error> { |
| 59 | + use std::path::PathBuf; |
| 60 | + let level = std::env::var("TURBOLITE_COMPRESSION_LEVEL") |
| 61 | + .ok() |
| 62 | + .and_then(|s| s.parse().ok()) |
| 63 | + .unwrap_or(3); |
| 64 | + let vfs = crate::CompressedVfs::new(PathBuf::from("."), level); |
| 65 | + crate::register("turbolite", vfs) |
| 66 | +} |
| 67 | + |
| 68 | +#[cfg(feature = "tiered")] |
| 69 | +fn register_tiered() -> Result<(), std::io::Error> { |
| 70 | + use std::path::PathBuf; |
| 71 | + use crate::tiered::{TieredConfig, TieredVfs}; |
| 72 | + |
| 73 | + let bucket = std::env::var("TURBOLITE_BUCKET") |
| 74 | + .expect("TURBOLITE_BUCKET must be set for tiered mode"); |
| 75 | + let prefix = std::env::var("TURBOLITE_PREFIX") |
| 76 | + .unwrap_or_else(|_| "turbolite".into()); |
| 77 | + let cache_dir = std::env::var("TURBOLITE_CACHE_DIR") |
| 78 | + .map(PathBuf::from) |
| 79 | + .unwrap_or_else(|_| PathBuf::from("/tmp/turbolite")); |
| 80 | + let endpoint_url = std::env::var("TURBOLITE_ENDPOINT_URL") |
| 81 | + .or_else(|_| std::env::var("AWS_ENDPOINT_URL")) |
| 82 | + .ok(); |
| 83 | + let region = std::env::var("TURBOLITE_REGION") |
| 84 | + .or_else(|_| std::env::var("AWS_REGION")) |
| 85 | + .ok(); |
| 86 | + let prefetch_threads = std::env::var("TURBOLITE_PREFETCH_THREADS") |
| 87 | + .ok() |
| 88 | + .and_then(|s| s.parse().ok()) |
| 89 | + .unwrap_or(0); // 0 = use default (num_cpus + 1) |
| 90 | + let compression_level = std::env::var("TURBOLITE_COMPRESSION_LEVEL") |
| 91 | + .ok() |
| 92 | + .and_then(|s| s.parse().ok()) |
| 93 | + .unwrap_or(3); |
| 94 | + let read_only = std::env::var("TURBOLITE_READ_ONLY") |
| 95 | + .map(|s| s == "1" || s == "true") |
| 96 | + .unwrap_or(false); |
| 97 | + |
| 98 | + let mut config = TieredConfig { |
| 99 | + bucket, |
| 100 | + prefix, |
| 101 | + cache_dir, |
| 102 | + endpoint_url, |
| 103 | + region, |
| 104 | + compression_level, |
| 105 | + read_only, |
| 106 | + ..Default::default() |
| 107 | + }; |
| 108 | + if prefetch_threads > 0 { |
| 109 | + config.prefetch_threads = prefetch_threads; |
| 110 | + } |
| 111 | + |
| 112 | + let vfs = TieredVfs::new(config)?; |
| 113 | + crate::tiered::register("turbolite", vfs) |
| 114 | +} |
| 115 | + |
| 116 | +#[cfg(not(feature = "tiered"))] |
| 117 | +fn register_tiered() -> Result<(), std::io::Error> { |
| 118 | + Err(std::io::Error::new( |
| 119 | + std::io::ErrorKind::Unsupported, |
| 120 | + "TURBOLITE_BUCKET is set but this extension was built without the 'tiered' feature", |
| 121 | + )) |
| 122 | +} |
0 commit comments