BE-930: Add Workspace::snapshot

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,
 };

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<SnapshotError>> {
+        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;

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 {

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<PathBuf> {
+        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")
+    }
 }

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<str>,
         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 `<target>.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 `<target>.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<SnapshotError>> {
+        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.