ursa-engine 0.7.0

Ursa's SQL engine session: a DataFusion session that reads files by path, locally or in object storage.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
//! The `.ursa` database file: a session's catalog, kept between processes.
//!
//! `ursa club.ursa` opens (or creates) a [`Database`] and installs it on the
//! session. The design is the Tempra page `/sql/ursa-file.md` in
//! `cl_dixon/ursa` (cldixon/ursa#237). On disk:
//!
//! ```text
//! club.ursa          the catalog: the statements that defined each object
//! club.ursa.d/       created on first use
//!   lock             held while a process has the database open for writing
//!   indexes/         stored graph indexes: `manifest.json` and `<hash>.idx`
//! ```
//!
//! This module owns the file, the sidecar directory and the lock. It stores
//! the catalog as an ordered list of statement texts ([`catalog_file`]) and
//! knows nothing of SQL: which statement defines which object, and replaying
//! them, is `ursa-sql`'s `database` module.
//!
//! A database is reached from the session config ([`database`],
//! [`database_of_config`]), so any code holding the session state or a task
//! context (the graph index cache, which stores indexes in it) can find it.
//!
//! **Relative paths** follow DuckDB. A relative file path in SQL (in `FROM`,
//! a graph's edge or vertex table, an external table's `LOCATION`, a
//! `COPY … TO` target) resolves against the current directory when the
//! statement runs, with or without a database. When a definition is stored in
//! the catalog, each of its local file paths is rewritten relative to the
//! database's directory ([`Database::stored_location`]), or made absolute if
//! it is outside that directory; replaying the catalog resolves the stored
//! paths against the database's directory ([`Database::replayed_location`]).
//! So a later process finds the same files from any directory, and a database
//! and its data files can move together. URLs are never rewritten.

use std::fs::File;
use std::path::{Component, Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};

use datafusion::error::{DataFusionError, Result};
use datafusion::execution::context::SessionState;
use datafusion::prelude::SessionConfig;

pub mod catalog_file;
pub mod index_file;
pub mod index_store;

pub use catalog_file::{CatalogFile, FORMAT, VERSION};
pub use index_file::IndexKey;
pub use index_store::IndexStore;

/// How long a writer waits for another process to close the database.
pub const DEFAULT_LOCK_TIMEOUT: Duration = Duration::from_secs(2);

/// How to open a database.
#[derive(Debug, Clone)]
pub struct OpenOptions {
    /// Open without the lock, and refuse statements that change the catalog.
    /// The file must exist.
    pub read_only: bool,
    /// How long to wait for the lock before failing with "… is in use by
    /// another process".
    pub lock_timeout: Duration,
}

impl Default for OpenOptions {
    fn default() -> Self {
        Self {
            read_only: false,
            lock_timeout: DEFAULT_LOCK_TIMEOUT,
        }
    }
}

impl OpenOptions {
    pub fn read_only() -> Self {
        Self {
            read_only: true,
            ..Self::default()
        }
    }
}

/// An open `.ursa` database.
#[derive(Debug)]
pub struct Database {
    /// The path as the user gave it: how messages name the database.
    name: String,
    /// The catalog file, absolute.
    path: PathBuf,
    /// The directory holding the catalog file, absolute: stored paths are
    /// relative to it.
    dir: PathBuf,
    /// `club.ursa.d`, absolute.
    sidecar: PathBuf,
    read_only: bool,
    /// The locked `club.ursa.d/lock`, held until the database is dropped.
    /// `None` when read-only.
    _lock: Option<File>,
    /// The catalog's statements, as last read or written.
    statements: Mutex<Vec<String>>,
    /// The stored graph indexes, opened on first use.
    indexes: OnceLock<Arc<IndexStore>>,
}

impl Database {
    /// Open the database at `path`, creating it (and its sidecar directory)
    /// unless read-only. A writer takes the lock first, waiting up to
    /// `options.lock_timeout`, then reads the catalog.
    pub fn open(path: impl AsRef<Path>, options: &OpenOptions) -> Result<Self> {
        let given = path.as_ref();
        let name = given.display().to_string();
        let path = std::path::absolute(given).map_err(|e| {
            DataFusionError::Execution(format!("cannot open database '{name}': {e}"))
        })?;
        let dir = path
            .parent()
            .map(Path::to_path_buf)
            .unwrap_or_else(|| PathBuf::from("/"));
        let file_name = path.file_name().ok_or_else(|| {
            DataFusionError::Execution(format!("'{name}' does not name a database file"))
        })?;
        let mut sidecar_name = file_name.to_os_string();
        sidecar_name.push(".d");
        let sidecar = dir.join(sidecar_name);

        if options.read_only {
            if !path.exists() {
                return Err(DataFusionError::Execution(format!(
                    "cannot open '{name}' read-only: the file does not exist"
                )));
            }
            let catalog = CatalogFile::read(&path, &name)?;
            return Ok(Self {
                name,
                path,
                dir,
                sidecar,
                read_only: true,
                _lock: None,
                statements: Mutex::new(catalog.statements),
                indexes: OnceLock::new(),
            });
        }

        if !dir.is_dir() {
            return Err(DataFusionError::Execution(format!(
                "cannot create database '{name}': the directory {} does not exist",
                dir.display()
            )));
        }
        let io = |what: &str, e: std::io::Error| {
            DataFusionError::Execution(format!("cannot open database '{name}': {what}: {e}"))
        };
        std::fs::create_dir_all(&sidecar)
            .map_err(|e| io(&format!("creating {}", sidecar.display()), e))?;
        let lock = File::options()
            .read(true)
            .write(true)
            .create(true)
            .truncate(false)
            .open(sidecar.join("lock"))
            .map_err(|e| io("opening its lock file", e))?;
        acquire(&lock, &name, options.lock_timeout)?;

        let statements = if path.exists() {
            CatalogFile::read(&path, &name)?.statements
        } else {
            let empty = CatalogFile::default();
            empty.write(&path, &sidecar, &name)?;
            empty.statements
        };
        Ok(Self {
            name,
            path,
            dir,
            sidecar,
            read_only: false,
            _lock: Some(lock),
            statements: Mutex::new(statements),
            indexes: OnceLock::new(),
        })
    }

    /// The database as the user named it (`club.ursa`).
    pub fn name(&self) -> &str {
        &self.name
    }

    /// The catalog file, absolute.
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// The directory holding the catalog file: stored relative paths are
    /// relative to it.
    pub fn dir(&self) -> &Path {
        &self.dir
    }

    /// The sidecar directory, `club.ursa.d`. It exists unless read-only.
    pub fn sidecar_dir(&self) -> &Path {
        &self.sidecar
    }

    /// Where stored graph indexes go (`club.ursa.d/indexes`), created on
    /// the first write.
    pub fn index_dir(&self) -> PathBuf {
        self.sidecar.join("indexes")
    }

    /// The database's stored graph indexes. A read-only database reads them
    /// and never writes.
    pub fn index_store(&self) -> Arc<IndexStore> {
        Arc::clone(
            self.indexes
                .get_or_init(|| Arc::new(IndexStore::new(self.index_dir(), !self.read_only))),
        )
    }

    /// Wait for stored-index writes still running.
    pub fn flush(&self) {
        if let Some(store) = self.indexes.get() {
            store.flush();
        }
    }

    pub fn is_read_only(&self) -> bool {
        self.read_only
    }

    /// The catalog's statements, in replay order.
    pub fn statements(&self) -> Vec<String> {
        self.statements
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .clone()
    }

    /// Replace the catalog with `statements`, rewriting the file atomically
    /// (temporary file, `fsync`, rename). Refused when read-only.
    pub fn save(&self, statements: Vec<String>) -> Result<()> {
        if self.read_only {
            return Err(DataFusionError::Plan(format!(
                "{} is open read-only",
                self.name
            )));
        }
        let mut current = self.statements.lock().unwrap_or_else(|e| e.into_inner());
        if *current == statements {
            return Ok(());
        }
        let file = CatalogFile {
            statements: statements.clone(),
        };
        file.write(&self.path, &self.sidecar, &self.name)?;
        *current = statements;
        Ok(())
    }

    /// How the catalog stores `location`, written in a statement run from
    /// `cwd`: a local path inside the database's directory becomes relative to
    /// it, any other local path becomes absolute. `None` when the location is
    /// stored as written: a URL, a `~/…` path, or a path already in its
    /// stored form.
    pub fn stored_location(&self, location: &str, cwd: &Path) -> Option<String> {
        if !is_local_path(location) {
            return None;
        }
        let absolute = normalize(&cwd.join(location));
        let canonical_dir = std::fs::canonicalize(&self.dir).ok();
        let inside = [Some(&self.dir), canonical_dir.as_ref()]
            .into_iter()
            .flatten()
            .find_map(|dir| absolute.strip_prefix(dir).ok());
        let mut stored = match inside {
            Some(relative) if relative.as_os_str().is_empty() => ".".to_string(),
            Some(relative) => relative.to_string_lossy().into_owned(),
            None => absolute.to_string_lossy().into_owned(),
        };
        keep_trailing_separator(location, &mut stored);
        (stored != location).then_some(stored)
    }

    /// `location` from a stored statement as replay runs it: a relative local
    /// path resolved against the database's directory; `None` for anything
    /// else (a URL, an absolute path, `~/…`), which runs as stored.
    pub fn replayed_location(&self, location: &str) -> Option<String> {
        if !is_relative_path(location) {
            return None;
        }
        let mut resolved = normalize(&self.dir.join(location))
            .to_string_lossy()
            .into_owned();
        keep_trailing_separator(location, &mut resolved);
        Some(resolved)
    }
}

impl Drop for Database {
    /// Finish writing stored indexes before the lock is released.
    fn drop(&mut self) {
        self.flush();
    }
}

/// Whether `location` is a local path (not a URL, not starting with `~`).
pub fn is_local_path(location: &str) -> bool {
    if location.is_empty() || location.starts_with('~') {
        return false;
    }
    match url::Url::parse(location) {
        // A one-letter scheme is a Windows drive letter.
        Ok(url) => url.scheme().len() == 1,
        Err(_) => true,
    }
}

/// Whether `location` is a relative local path (not a URL, not absolute, not
/// starting with `~`).
pub fn is_relative_path(location: &str) -> bool {
    is_local_path(location) && Path::new(location).is_relative()
}

/// `path` with `.` and `..` components removed lexically, without touching
/// the filesystem (the path may be a glob, or not exist yet).
fn normalize(path: &Path) -> PathBuf {
    let mut out = PathBuf::new();
    for component in path.components() {
        match component {
            Component::CurDir => {}
            Component::ParentDir => match out.components().next_back() {
                Some(Component::Normal(_)) => {
                    out.pop();
                }
                // `..` at the root stays at the root.
                Some(Component::RootDir | Component::Prefix(_)) => {}
                _ => out.push(".."),
            },
            other => out.push(other),
        }
    }
    out
}

/// A directory written with a trailing separator (`data/`) keeps it.
fn keep_trailing_separator(original: &str, rewritten: &mut String) {
    let separator = std::path::MAIN_SEPARATOR;
    if (original.ends_with('/') || original.ends_with(separator))
        && !rewritten.ends_with(['/', separator])
    {
        rewritten.push(separator);
    }
}

/// Take the exclusive lock on `lock`, retrying until `timeout` has passed.
fn acquire(lock: &File, name: &str, timeout: Duration) -> Result<()> {
    let start = Instant::now();
    loop {
        match lock.try_lock() {
            Ok(()) => return Ok(()),
            Err(std::fs::TryLockError::WouldBlock) => {
                if start.elapsed() >= timeout {
                    return Err(DataFusionError::Execution(format!(
                        "{name} is in use by another process"
                    )));
                }
                std::thread::sleep(Duration::from_millis(50));
            }
            Err(std::fs::TryLockError::Error(e)) => {
                return Err(DataFusionError::Execution(format!(
                    "cannot lock database '{name}': {e}"
                )))
            }
        }
    }
}

/// Where a session keeps its database: empty for an in-memory session, set
/// once by opening or attaching one. Installed on every session's config by
/// `new_session`, so `ATTACH` can fill it later.
#[derive(Debug, Default)]
pub struct DatabaseSlot {
    database: OnceLock<Arc<Database>>,
}

impl DatabaseSlot {
    pub fn get(&self) -> Option<Arc<Database>> {
        self.database.get().cloned()
    }

    /// Put `database` in the slot; fails if the session already has one.
    pub fn set(&self, database: Arc<Database>) -> Result<()> {
        self.database
            .set(database)
            .map_err(|_| DataFusionError::Plan("this session already has a database open".into()))
    }
}

/// Install an empty database slot on `config` (done by `new_session`).
pub(crate) fn with_database_slot(config: SessionConfig) -> SessionConfig {
    config.with_extension(Arc::new(DatabaseSlot::default()))
}

/// The session's database slot.
pub fn slot(state: &SessionState) -> Arc<DatabaseSlot> {
    slot_of_config(state.config())
}

fn slot_of_config(config: &SessionConfig) -> Arc<DatabaseSlot> {
    config
        .get_extension::<DatabaseSlot>()
        .expect("an Ursa session has a database slot (built by ursa_engine::new_session)")
}

/// The session's database, if it has one.
pub fn database(state: &SessionState) -> Option<Arc<Database>> {
    database_of_config(state.config())
}

/// The database of the session `config` belongs to, if it has one. A task
/// context's `session_config()` reaches it too.
pub fn database_of_config(config: &SessionConfig) -> Option<Arc<Database>> {
    config.get_extension::<DatabaseSlot>()?.get()
}

/// Wait for the session's stored-index writes, if it has a database. Call it
/// before a process exits, so an index built by its last query is stored.
pub fn flush(state: &SessionState) {
    if let Some(db) = database(state) {
        db.flush();
    }
}

#[cfg(test)]
mod tests;