diff --git a/libs/@local/integrations/engine/src/lib.rs b/libs/@local/integrations/engine/src/lib.rs index f9ab7aa3ad3..ff4376b05ca 100644 --- a/libs/@local/integrations/engine/src/lib.rs +++ b/libs/@local/integrations/engine/src/lib.rs @@ -20,6 +20,6 @@ mod variant_name; mod workspace; pub use self::workspace::{ - CloseError, DataErrorType, OpenError, QueryErrorType, Rows, Setting, StatementError, Workspace, - WorkspaceConfig, + CloseError, DataErrorType, OpenError, QueryErrorType, Rows, Setting, SnapshotError, + StatementError, Workspace, WorkspaceConfig, }; diff --git a/libs/@local/integrations/engine/src/workspace/connection.rs b/libs/@local/integrations/engine/src/workspace/connection.rs index e824118df95..1a13a9738da 100644 --- a/libs/@local/integrations/engine/src/workspace/connection.rs +++ b/libs/@local/integrations/engine/src/workspace/connection.rs @@ -1,9 +1,9 @@ use alloc::sync::Arc; use core::{error::Error, panic::AssertUnwindSafe}; use std::{ - fs::{File, OpenOptions, TryLockError}, - panic, - path::PathBuf, + fs::{self, File, OpenOptions, TryLockError}, + io, panic, + path::{Component, Path, PathBuf}, }; use duckdb::{Config, types::Value}; @@ -13,7 +13,7 @@ use tokio::sync::mpsc; use crate::{ sql::Identifier, workspace::{ - CloseError, Command, OpenError, Rows, Setting, StatementError, + CloseError, Command, OpenError, Rows, Setting, SnapshotError, StatementError, files::WorkspaceFiles, interrupt::{Interrupt, InterruptGuard, RequestId}, reply::Reply, @@ -28,6 +28,7 @@ use crate::{ pub(super) struct Connection { interrupt: InterruptGuard, duckdb: duckdb::Connection, + files: WorkspaceFiles, /// An exclusive lock on the workspace's lock file, held until the database is closed. lock: File, } @@ -87,6 +88,7 @@ impl Connection { Ok(Self { interrupt: InterruptGuard::new(Arc::new(Interrupt::new(duckdb.interrupt_handle()))), duckdb, + files, lock, }) } @@ -163,6 +165,70 @@ impl Connection { Ok(()) } + /// Copies the database file to `partial`, syncs it and renames it to `target`. + /// + /// An existing file at `partial` is removed first, so the copy cannot write through a link to + /// another file, such as the database file itself. + fn write_copy(&self, partial: &Path, target: &Path) -> io::Result<()> { + match fs::remove_file(partial) { + Ok(()) => {} + Err(error) if error.kind() == io::ErrorKind::NotFound => {} + Err(error) => return Err(error), + } + let mut copy = OpenOptions::new() + .write(true) + .create_new(true) + .open(partial)?; + io::copy(&mut File::open(self.files.database())?, &mut copy)?; + copy.sync_all()?; + drop(copy); + fs::rename(partial, target)?; + // The engine chooses `target`, and `snapshot` refuses relative paths and `..` before this + // runs. Opening the directory only syncs the rename to disk. + // nosemgrep: rust.actix.path-traversal.tainted-path.tainted-path + File::open(WorkspaceFiles::directory(target))?.sync_all() + } + + #[tracing::instrument(level = "debug", skip(self))] + fn snapshot(&self, target: &Path) -> Result<(), Report> { + if !target.is_absolute() + || target + .components() + .any(|part| matches!(part, Component::ParentDir)) + { + return Err(Report::new(SnapshotError::InvalidTarget { + target: target.to_owned(), + })); + } + let partial = WorkspaceFiles::partial(target); + if self.files.contains(target) || self.files.contains(&partial) { + return Err(Report::new(SnapshotError::WorkspaceFile { + target: target.to_owned(), + })); + } + + // No statement can run between the checkpoint and the copy, because this thread owns the + // only connection. + self.duckdb + .execute_batch("CHECKPOINT") + .change_context(SnapshotError::Checkpoint)?; + + self.write_copy(&partial, target).map_err(|error| { + if let Err(removal) = fs::remove_file(&partial) + && removal.kind() != io::ErrorKind::NotFound + { + tracing::error!( + error = ?removal, + path = %partial.display(), + "could not remove the partial snapshot" + ); + } + Report::new(error).change_context(SnapshotError::Write { + target: target.to_owned(), + }) + }) + } + /// Checkpoints the database, then closes it whether or not the checkpoint succeeded. /// /// DuckDB also checkpoints when it closes the database, but it cannot report a failure from @@ -219,6 +285,9 @@ impl Connection { connection.create_table(&table, &sql) }); } + Command::Snapshot { target, reply } => { + self.answer(id, reply, |connection| connection.snapshot(&target)); + } Command::Close { reply } => { reply.send(self.close()); return; diff --git a/libs/@local/integrations/engine/src/workspace/error.rs b/libs/@local/integrations/engine/src/workspace/error.rs index 1100eb1f84f..becec93c3e2 100644 --- a/libs/@local/integrations/engine/src/workspace/error.rs +++ b/libs/@local/integrations/engine/src/workspace/error.rs @@ -214,6 +214,21 @@ impl StatementError { } } +/// Reports why [`Workspace::snapshot`](crate::Workspace::snapshot) failed. +#[derive(Debug, derive_more::Display, derive_more::Error)] +pub enum SnapshotError { + #[display("`{}` is relative or contains `..`", target.display())] + InvalidTarget { target: PathBuf }, + #[display("`{}` is one of the workspace's own files", target.display())] + WorkspaceFile { target: PathBuf }, + #[display("could not checkpoint the database")] + Checkpoint, + #[display("could not write the snapshot to `{}`", target.display())] + Write { target: PathBuf }, + #[display("the workspace thread stopped before the snapshot finished")] + Stopped, +} + /// Reports why [`Workspace::close`](crate::Workspace::close) failed. #[derive(Debug, derive_more::Display, derive_more::Error)] pub enum CloseError { diff --git a/libs/@local/integrations/engine/src/workspace/files.rs b/libs/@local/integrations/engine/src/workspace/files.rs index a5efba4483e..a62a6cc372c 100644 --- a/libs/@local/integrations/engine/src/workspace/files.rs +++ b/libs/@local/integrations/engine/src/workspace/files.rs @@ -1,19 +1,54 @@ -use std::path::{Path, PathBuf}; +use std::{ + fs, + path::{Path, PathBuf}, +}; /// The files of a workspace, named after its database file. pub(super) struct WorkspaceFiles { database: PathBuf, + /// DuckDB's write-ahead log. + log: PathBuf, + /// DuckDB's default temporary directory. + temp: PathBuf, /// The file that the workspace locks while the database is open. lock: PathBuf, } impl WorkspaceFiles { + /// Returns `path` with `suffix` appended to its file name. + fn with_suffix(path: &Path, suffix: &str) -> PathBuf { + let mut name = path.as_os_str().to_owned(); + name.push(suffix); + PathBuf::from(name) + } + + /// Returns the directory that holds `path`, which is the current directory for a bare file + /// name. + pub(super) fn directory(path: &Path) -> &Path { + path.parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")) + } + + /// Returns `path` with its directory resolved, or `None` if the directory cannot be resolved + /// or `path` has no file name. + fn resolve(path: &Path) -> Option { + Some( + fs::canonicalize(Self::directory(path)) + .ok()? + .join(path.file_name()?), + ) + } + + /// Returns the files of the workspace whose database file is `database`. + /// + /// DuckDB names its log and its default temporary directory after the database file. pub(super) fn new(database: PathBuf) -> Self { - let mut lock = database.as_os_str().to_owned(); - lock.push(".lock"); Self { + log: Self::with_suffix(&database, ".wal"), + temp: Self::with_suffix(&database, ".tmp"), + lock: Self::with_suffix(&database, ".lock"), database, - lock: PathBuf::from(lock), } } @@ -24,4 +59,20 @@ impl WorkspaceFiles { pub(super) fn lock(&self) -> &Path { &self.lock } + + /// Returns whether `path` names one of the files, once both directories are resolved. + pub(super) fn contains(&self, path: &Path) -> bool { + let Some(path) = Self::resolve(path) else { + return false; + }; + [&self.database, &self.log, &self.temp, &self.lock] + .into_iter() + .any(|file| Self::resolve(file).is_some_and(|file| file == path)) + } + + /// Returns `target` with `.partial` appended, the file a copy is written to before it is + /// renamed to `target`. + pub(super) fn partial(target: &Path) -> PathBuf { + Self::with_suffix(target, ".partial") + } } diff --git a/libs/@local/integrations/engine/src/workspace/mod.rs b/libs/@local/integrations/engine/src/workspace/mod.rs index 1563f454be2..fe68e4ebde1 100644 --- a/libs/@local/integrations/engine/src/workspace/mod.rs +++ b/libs/@local/integrations/engine/src/workspace/mod.rs @@ -14,7 +14,10 @@ use error_stack::{Report, ResultExt as _}; use tokio::sync::mpsc; pub use self::{ - error::{CloseError, DataErrorType, OpenError, QueryErrorType, Setting, StatementError}, + error::{ + CloseError, DataErrorType, OpenError, QueryErrorType, Setting, SnapshotError, + StatementError, + }, rows::Rows, }; use crate::workspace::{ @@ -51,6 +54,10 @@ enum Command { sql: Box, reply: Reply<(), StatementError>, }, + Snapshot { + target: PathBuf, + reply: Reply<(), SnapshotError>, + }, Close { reply: Reply<(), CloseError>, }, @@ -201,6 +208,34 @@ impl Workspace { .await } + /// Writes a copy of the database to `target` that DuckDB can open on its own. + /// + /// The copy holds the tables of every request sent before the snapshot, and none sent after + /// it. It is written to a new file at `.partial`, which replaces any file there, synced + /// to disk, and then renamed to `target`, so it replaces an existing file at `target` only once + /// it is complete. + /// + /// # Errors + /// + /// - [`InvalidTarget`] if `target` is relative or contains `..` + /// - [`WorkspaceFile`] if `target` or `.partial` is one of the workspace's files + /// - [`Checkpoint`] if DuckDB cannot write its log into the database file + /// - [`Write`] if the copy cannot be written, synced or renamed + /// - [`Stopped`] if the workspace thread stopped + /// + /// [`InvalidTarget`]: SnapshotError::InvalidTarget + /// [`WorkspaceFile`]: SnapshotError::WorkspaceFile + /// [`Checkpoint`]: SnapshotError::Checkpoint + /// [`Write`]: SnapshotError::Write + /// [`Stopped`]: SnapshotError::Stopped + pub async fn snapshot(&self, target: PathBuf) -> Result<(), Report> { + self.request( + |reply| Command::Snapshot { target, reply }, + SnapshotError::Stopped, + ) + .await + } + /// Checkpoints and closes the database after the requests already sent have run. /// /// The database is closed even if the checkpoint fails. diff --git a/libs/@local/integrations/engine/tests/workspace.rs b/libs/@local/integrations/engine/tests/workspace.rs index ee0560386f5..09234840a7b 100644 --- a/libs/@local/integrations/engine/tests/workspace.rs +++ b/libs/@local/integrations/engine/tests/workspace.rs @@ -1,12 +1,16 @@ #![expect(missing_docs, reason = "integration tests")] use core::{future::poll_fn, num::NonZeroUsize, pin::Pin, task::Poll, time::Duration}; -use std::path::PathBuf; +use std::{ + fs, + path::{Path, PathBuf}, +}; use duckdb::types::{TimeUnit, Value}; use error_stack::Report; use hash_integrations_engine::{ - DataErrorType, OpenError, QueryErrorType, StatementError, Workspace, WorkspaceConfig, + DataErrorType, OpenError, QueryErrorType, SnapshotError, StatementError, Workspace, + WorkspaceConfig, }; use pretty_assertions::assert_eq; use tempfile::TempDir; @@ -27,6 +31,16 @@ async fn first_value(workspace: &Workspace, sql: &str) -> Result i64 { + duckdb::Connection::open(path) + .expect("the database should open") + .query_row(&format!("SELECT count(*) FROM {table}"), [], |row| { + row.get(0) + }) + .expect("the database should hold the table") +} + /// Polls `request` once and checks that it is still pending. async fn send(mut request: Pin<&mut impl Future>) { assert!( @@ -443,3 +457,143 @@ async fn dropping_workspace_interrupts_running_query() { .expect("the lock should be released once the dropped workspace interrupts its query"); reopened.expect("the database should open once the lock is released"); } + +#[tokio::test] +async fn snapshot_holds_tables_created_before_it() { + let directory = TempDir::new().expect("should create a directory"); + let target = directory.path().join("snapshot.duckdb"); + fs::write(&target, "not a database").expect("should write a file"); + let workspace = open(&directory).await; + + workspace + .create_table("flights".into(), "SELECT 'BA117' AS code".into()) + .await + .expect("the table should be created"); + workspace + .snapshot(target.clone()) + .await + .expect("the snapshot should be written"); + workspace + .create_table( + "flights".into(), + "SELECT * FROM (VALUES ('BA117'), ('LH400')) AS flights (code)".into(), + ) + .await + .expect("the table should be replaced"); + + assert_eq!( + count_rows(&target, "flights"), + 1, + "the snapshot should hold only the first table" + ); + assert!( + !directory.path().join("snapshot.duckdb.partial").exists(), + "the partial copy should be renamed to the target" + ); +} + +#[tokio::test] +async fn snapshot_refuses_invalid_targets() { + let directory = TempDir::new().expect("should create a directory"); + fs::create_dir_all(directory.path().join("nested")).expect("should create a directory"); + let workspace = open(&directory).await; + + for target in [ + PathBuf::from("snapshot.duckdb"), + directory.path().join("nested/../snapshot.duckdb"), + ] { + let report = workspace + .snapshot(target.clone()) + .await + .expect_err("a relative target or one with `..` should be refused"); + assert!( + matches!( + report.current_context(), + SnapshotError::InvalidTarget { .. } + ), + "`{}` should be refused, not by {report:?}", + target.display() + ); + } +} + +#[tokio::test] +async fn snapshot_refuses_workspace_files() { + let directory = TempDir::new().expect("should create a directory"); + let workspace = open(&directory).await; + + for target in [ + database_path(&directory), + directory.path().join("workspace.duckdb.wal"), + directory.path().join("workspace.duckdb.tmp"), + directory.path().join("workspace.duckdb.lock"), + ] { + let report = workspace + .snapshot(target.clone()) + .await + .expect_err("a workspace file should not be a snapshot target"); + assert!( + matches!( + report.current_context(), + SnapshotError::WorkspaceFile { .. } + ), + "`{}` should be refused, not by {report:?}", + target.display() + ); + } +} + +#[tokio::test] +async fn snapshot_failure_removes_partial_copy() { + let directory = TempDir::new().expect("should create a directory"); + let target = directory.path().join("snapshot"); + fs::create_dir_all(&target).expect("should create a directory"); + let workspace = open(&directory).await; + + let report = workspace + .snapshot(target) + .await + .expect_err("a directory should not be replaced by the snapshot"); + + assert!( + matches!(report.current_context(), SnapshotError::Write { .. }), + "the failed write should be reported, not {report:?}" + ); + assert!( + !directory.path().join("snapshot.partial").exists(), + "the partial copy should be removed" + ); +} + +#[tokio::test] +async fn snapshot_replaces_link_at_partial_path() { + let directory = TempDir::new().expect("should create a directory"); + let workspace = open(&directory).await; + workspace + .create_table("flights".into(), "SELECT 'BA117' AS code".into()) + .await + .expect("the table should be created"); + + let target = directory.path().join("snapshot"); + fs::hard_link( + database_path(&directory), + directory.path().join("snapshot.partial"), + ) + .expect("should link the database file"); + workspace + .snapshot(target.clone()) + .await + .expect("the snapshot should replace the link"); + workspace.close().await.expect("the workspace should close"); + + assert_eq!( + count_rows(&database_path(&directory), "flights"), + 1, + "the database should keep its table" + ); + assert_eq!( + count_rows(&target, "flights"), + 1, + "the snapshot should hold the table" + ); +}