1use std::collections::BTreeMap;
35use std::fs::File;
36use std::io::{Read, Seek, Write};
37use std::path::{Component, Path, PathBuf};
38use std::sync::{Arc, Mutex, OnceLock};
39use std::time::{Duration, Instant};
40
41use datafusion::error::{DataFusionError, Result};
42use datafusion::execution::context::SessionState;
43use datafusion::prelude::SessionConfig;
44
45pub mod catalog_file;
46pub mod index_file;
47pub mod index_store;
48
49use crate::options::{ByteSize, DEFAULT_INDEX_STORAGE_LIMIT};
50
51pub use catalog_file::{CatalogFile, FORMAT, VERSION};
52pub use index_file::IndexKey;
53pub use index_store::IndexStore;
54
55pub const DEFAULT_LOCK_TIMEOUT: Duration = Duration::from_secs(2);
57
58#[derive(Debug, Clone)]
60pub struct OpenOptions {
61 pub read_only: bool,
64 pub lock_timeout: Duration,
67}
68
69impl Default for OpenOptions {
70 fn default() -> Self {
71 Self {
72 read_only: false,
73 lock_timeout: DEFAULT_LOCK_TIMEOUT,
74 }
75 }
76}
77
78impl OpenOptions {
79 pub fn read_only() -> Self {
80 Self {
81 read_only: true,
82 ..Self::default()
83 }
84 }
85}
86
87#[derive(Debug)]
89pub struct Database {
90 name: String,
92 path: PathBuf,
94 dir: PathBuf,
97 sidecar: PathBuf,
99 read_only: bool,
100 _lock: Option<File>,
103 catalog: Mutex<CatalogFile>,
105 indexes: OnceLock<Arc<IndexStore>>,
107}
108
109pub const INDEX_STORAGE_LIMIT: &str = "ursa.index_storage_limit";
111
112impl Database {
113 pub fn open(path: impl AsRef<Path>, options: &OpenOptions) -> Result<Self> {
117 let given = path.as_ref();
118 let name = given.display().to_string();
119 let path = std::path::absolute(given).map_err(|e| {
120 DataFusionError::Execution(format!("cannot open database '{name}': {e}"))
121 })?;
122 let dir = path
123 .parent()
124 .map(Path::to_path_buf)
125 .unwrap_or_else(|| PathBuf::from("/"));
126 let file_name = path.file_name().ok_or_else(|| {
127 DataFusionError::Execution(format!("'{name}' does not name a database file"))
128 })?;
129 let mut sidecar_name = file_name.to_os_string();
130 sidecar_name.push(".d");
131 let sidecar = dir.join(sidecar_name);
132
133 if options.read_only {
134 if !path.exists() {
135 return Err(DataFusionError::Execution(format!(
136 "cannot open '{name}' read-only: the file does not exist"
137 )));
138 }
139 let catalog = CatalogFile::read(&path, &name)?;
140 return Ok(Self {
141 name,
142 path,
143 dir,
144 sidecar,
145 read_only: true,
146 _lock: None,
147 catalog: Mutex::new(catalog),
148 indexes: OnceLock::new(),
149 });
150 }
151
152 if !dir.is_dir() {
153 return Err(DataFusionError::Execution(format!(
154 "cannot create database '{name}': the directory {} does not exist",
155 dir.display()
156 )));
157 }
158 let io = |what: &str, e: std::io::Error| {
159 DataFusionError::Execution(format!("cannot open database '{name}': {what}: {e}"))
160 };
161 std::fs::create_dir_all(&sidecar)
162 .map_err(|e| io(&format!("creating {}", sidecar.display()), e))?;
163 let mut lock = File::options()
164 .read(true)
165 .write(true)
166 .create(true)
167 .truncate(false)
168 .open(sidecar.join("lock"))
169 .map_err(|e| io("opening its lock file", e))?;
170 acquire(&lock, &name, options.lock_timeout)?;
171 record_holder(&mut lock);
172
173 let catalog = if path.exists() {
174 CatalogFile::read(&path, &name)?
175 } else {
176 let empty = CatalogFile::default();
177 empty.write(&path, &sidecar, &name)?;
178 empty
179 };
180 Ok(Self {
181 name,
182 path,
183 dir,
184 sidecar,
185 read_only: false,
186 _lock: Some(lock),
187 catalog: Mutex::new(catalog),
188 indexes: OnceLock::new(),
189 })
190 }
191
192 pub fn name(&self) -> &str {
194 &self.name
195 }
196
197 pub fn path(&self) -> &Path {
199 &self.path
200 }
201
202 pub fn dir(&self) -> &Path {
205 &self.dir
206 }
207
208 pub fn sidecar_dir(&self) -> &Path {
210 &self.sidecar
211 }
212
213 pub fn index_dir(&self) -> PathBuf {
216 self.sidecar.join("indexes")
217 }
218
219 pub fn index_store(&self) -> Arc<IndexStore> {
222 Arc::clone(self.indexes.get_or_init(|| {
223 Arc::new(IndexStore::new(
224 self.index_dir(),
225 !self.read_only,
226 self.name.clone(),
227 self.index_storage_limit(),
228 ))
229 }))
230 }
231
232 pub fn flush(&self) {
234 if let Some(store) = self.indexes.get() {
235 store.flush();
236 }
237 }
238
239 pub fn is_read_only(&self) -> bool {
240 self.read_only
241 }
242
243 fn catalog(&self) -> std::sync::MutexGuard<'_, CatalogFile> {
244 self.catalog.lock().unwrap_or_else(|e| e.into_inner())
245 }
246
247 pub fn statements(&self) -> Vec<String> {
249 self.catalog().statements.clone()
250 }
251
252 pub fn settings(&self) -> BTreeMap<String, String> {
254 self.catalog().settings.clone()
255 }
256
257 pub fn index_storage_limit(&self) -> u64 {
261 self.catalog()
262 .settings
263 .get(INDEX_STORAGE_LIMIT)
264 .and_then(|v| ByteSize::parse(v).ok())
265 .map_or(DEFAULT_INDEX_STORAGE_LIMIT, |b| b.0)
266 }
267
268 pub fn save(&self, statements: Vec<String>) -> Result<()> {
272 let mut catalog = self.catalog();
273 if !self.read_only && catalog.statements == statements {
274 return Ok(());
275 }
276 let updated = CatalogFile {
277 statements,
278 settings: catalog.settings.clone(),
279 };
280 self.write_catalog(&mut catalog, updated)
281 }
282
283 pub fn set_setting(&self, name: &str, value: Option<&str>) -> Result<()> {
287 if name == INDEX_STORAGE_LIMIT {
288 if let Some(value) = value {
289 ByteSize::parse(value)?;
290 }
291 }
292 let mut catalog = self.catalog();
293 let mut updated = catalog.clone();
294 match value {
295 Some(value) => updated.settings.insert(name.to_string(), value.to_string()),
296 None => updated.settings.remove(name),
297 };
298 if !self.read_only && updated == *catalog {
299 return Ok(());
300 }
301 self.write_catalog(&mut catalog, updated)?;
302 drop(catalog);
303 if let Some(store) = self.indexes.get() {
304 store.set_limit(self.index_storage_limit());
305 }
306 Ok(())
307 }
308
309 fn write_catalog(&self, current: &mut CatalogFile, updated: CatalogFile) -> Result<()> {
310 if self.read_only {
311 return Err(DataFusionError::Plan(format!(
312 "{} is open read-only",
313 self.name
314 )));
315 }
316 updated.write(&self.path, &self.sidecar, &self.name)?;
317 *current = updated;
318 Ok(())
319 }
320
321 pub fn stored_location(&self, location: &str, cwd: &Path) -> Option<String> {
327 if !is_local_path(location) {
328 return None;
329 }
330 let absolute = normalize(&cwd.join(location));
331 let canonical_dir = std::fs::canonicalize(&self.dir).ok();
332 let inside = [Some(&self.dir), canonical_dir.as_ref()]
333 .into_iter()
334 .flatten()
335 .find_map(|dir| absolute.strip_prefix(dir).ok());
336 let mut stored = match inside {
337 Some(relative) if relative.as_os_str().is_empty() => ".".to_string(),
338 Some(relative) => relative.to_string_lossy().into_owned(),
339 None => absolute.to_string_lossy().into_owned(),
340 };
341 keep_trailing_separator(location, &mut stored);
342 (stored != location).then_some(stored)
343 }
344
345 pub fn replayed_location(&self, location: &str) -> Option<String> {
349 if !is_relative_path(location) {
350 return None;
351 }
352 let mut resolved = normalize(&self.dir.join(location))
353 .to_string_lossy()
354 .into_owned();
355 keep_trailing_separator(location, &mut resolved);
356 Some(resolved)
357 }
358}
359
360impl Drop for Database {
361 fn drop(&mut self) {
363 self.flush();
364 }
365}
366
367pub fn is_local_path(location: &str) -> bool {
369 if location.is_empty() || location.starts_with('~') {
370 return false;
371 }
372 match url::Url::parse(location) {
373 Ok(url) => url.scheme().len() == 1,
375 Err(_) => true,
376 }
377}
378
379pub fn is_relative_path(location: &str) -> bool {
382 is_local_path(location) && Path::new(location).is_relative()
383}
384
385fn normalize(path: &Path) -> PathBuf {
388 let mut out = PathBuf::new();
389 for component in path.components() {
390 match component {
391 Component::CurDir => {}
392 Component::ParentDir => match out.components().next_back() {
393 Some(Component::Normal(_)) => {
394 out.pop();
395 }
396 Some(Component::RootDir | Component::Prefix(_)) => {}
398 _ => out.push(".."),
399 },
400 other => out.push(other),
401 }
402 }
403 out
404}
405
406fn keep_trailing_separator(original: &str, rewritten: &mut String) {
408 let separator = std::path::MAIN_SEPARATOR;
409 if (original.ends_with('/') || original.ends_with(separator))
410 && !rewritten.ends_with(['/', separator])
411 {
412 rewritten.push(separator);
413 }
414}
415
416fn acquire(lock: &File, name: &str, timeout: Duration) -> Result<()> {
418 let start = Instant::now();
419 loop {
420 match lock.try_lock() {
421 Ok(()) => return Ok(()),
422 Err(std::fs::TryLockError::WouldBlock) => {
423 if start.elapsed() >= timeout {
424 let holder = holder(lock)
425 .map(|pid| format!(" (process {pid})"))
426 .unwrap_or_default();
427 return Err(DataFusionError::Execution(format!(
428 "{name} is in use by another process{holder}, which has it open for \
429 writing; open it read-only to query it meanwhile (ursa --readonly \
430 {name}), or wait for that process to finish"
431 )));
432 }
433 std::thread::sleep(Duration::from_millis(50));
434 }
435 Err(std::fs::TryLockError::Error(e)) => {
436 return Err(DataFusionError::Execution(format!(
437 "cannot lock database '{name}': {e}"
438 )))
439 }
440 }
441 }
442}
443
444fn record_holder(lock: &mut File) {
447 let _ = lock.set_len(0).and_then(|()| {
448 lock.rewind()?;
449 writeln!(lock, "{}", std::process::id())?;
450 lock.flush()
451 });
452}
453
454fn holder(lock: &File) -> Option<u32> {
456 let mut text = String::new();
457 let mut reader = lock;
458 reader.rewind().ok()?;
459 reader.take(32).read_to_string(&mut text).ok()?;
460 text.trim().parse().ok()
461}
462
463#[derive(Debug, Default)]
467pub struct DatabaseSlot {
468 database: OnceLock<Arc<Database>>,
469}
470
471impl DatabaseSlot {
472 pub fn get(&self) -> Option<Arc<Database>> {
473 self.database.get().cloned()
474 }
475
476 pub fn set(&self, database: Arc<Database>) -> Result<()> {
478 self.database
479 .set(database)
480 .map_err(|_| DataFusionError::Plan("this session already has a database open".into()))
481 }
482}
483
484pub(crate) fn with_database_slot(config: SessionConfig) -> SessionConfig {
486 config.with_extension(Arc::new(DatabaseSlot::default()))
487}
488
489pub fn slot(state: &SessionState) -> Arc<DatabaseSlot> {
491 slot_of_config(state.config())
492}
493
494fn slot_of_config(config: &SessionConfig) -> Arc<DatabaseSlot> {
495 config
496 .get_extension::<DatabaseSlot>()
497 .expect("an Ursa session has a database slot (built by ursa_engine::new_session)")
498}
499
500pub fn database(state: &SessionState) -> Option<Arc<Database>> {
502 database_of_config(state.config())
503}
504
505pub fn database_of_config(config: &SessionConfig) -> Option<Arc<Database>> {
508 config.get_extension::<DatabaseSlot>()?.get()
509}
510
511pub fn flush(state: &SessionState) {
514 if let Some(db) = database(state) {
515 db.flush();
516 }
517}
518
519#[cfg(test)]
520mod tests;