Initial commit

.gitignore

@@ -0,0 +1,2 @@
+/target
+Cargo.lock

Cargo.toml

@@ -0,0 +1,33 @@
+[package]
+name = "tmp-postgrust"
+version = "0.5.0"
+authors = ["John Children <[email protected]>"]
+license = "MIT"
+edition = "2018"
+description = "Temporary postgresql instances for testing"
+repository = "https://github.com/CQCL/tmp-postgrust"
+readme = "README.md"
+keywords = ["testing", "database", "postgres"]
+
+[badges]
+maintenance = { status = "experimental" }
+
+[dependencies]
+glob = "0.3"
+lazy_static = "1.4.0"
+nix = "0.22"
+tempdir = "0.3"
+thiserror = "1.0"
+tokio = { version = "1.8", features = ["parking_lot", "rt", "sync", "io-util", "process", "macros", "fs"], default-features = false, optional = true }
+tracing = "0.1"
+which = "4.0"
+
+[dev-dependencies]
+test-env-log = { version = "0.2", default-features = false, features = ["trace"] }
+tokio = { version = "1.8", features = ["parking_lot", "rt", "rt-multi-thread", "sync", "io-util", "process", "macros", "fs"], default-features = false }
+tokio-postgres = "0.7"
+tracing-subscriber = { version = "0.2", default-features = false, features = ["env-filter", "fmt"] }
+
+[features]
+default = []
+tokio-process = ["tokio"]

LICENSE

@@ -0,0 +1,25 @@
+Copyright (c) 2021 Cambridge Quantum Computing
+
+Permission is hereby granted, free of charge, to any
+person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the
+Software without restriction, including without
+limitation the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software
+is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice
+shall be included in all copies or substantial portions
+of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.

README.md

@@ -0,0 +1,13 @@
+![Maintenance](https://img.shields.io/badge/maintenance-experimental-blue.svg)
+
+# tmp-postgrust
+
+`tmp-postgrust` provides temporary postgresql processes that are cleaned up
+after being dropped.
+
+
+## Inspiration / Similar Projects
+- [tmp-postgres](https://github.com/jfischoff/tmp-postgres)
+- [testing.postgresql](https://github.com/tk0miya/testing.postgresql)
+
+License: MIT

src/asynchronous.rs

@@ -0,0 +1,190 @@
+use std::path::Path;
+use std::process::Stdio;
+use std::sync::Arc;
+
+use tempdir::TempDir;
+use tokio::io::Lines;
+use tokio::process::{ChildStderr, ChildStdout};
+
+use tokio::sync::oneshot::Sender;
+use tokio::sync::{Semaphore, SemaphorePermit};
+use tokio::{
+    io::BufReader,
+    process::{Child, Command},
+};
+use tracing::{debug, instrument};
+
+use crate::errors::{ProcessCapture, TmpPostgrustError, TmpPostgrustResult};
+use crate::search::find_postgresql_command;
+
+/// Limit the total processes that can be running at any one time.
+pub(crate) static MAX_CONCURRENT_PROCESSES: Semaphore = Semaphore::const_new(8);
+
+#[instrument(skip(command, fail))]
+async fn exec_process(
+    command: &mut Command,
+    fail: impl FnOnce(ProcessCapture) -> TmpPostgrustError,
+) -> TmpPostgrustResult<()> {
+    debug!("running command: {:?}", command);
+
+    let output = command
+        .output()
+        .await
+        .map_err(|err| TmpPostgrustError::ExecSubprocessFailed {
+            source: err,
+            command: format!("{:?}", command),
+        })?;
+
+    if output.status.success() {
+        for line in String::from_utf8(output.stdout).unwrap().lines() {
+            debug!("{}", line);
+        }
+        Ok(())
+    } else {
+        Err(fail(ProcessCapture {
+            stdout: String::from_utf8(output.stdout).unwrap(),
+            stderr: String::from_utf8(output.stderr).unwrap(),
+        }))
+    }
+}
+
+#[instrument]
+pub(crate) fn start_postgres_subprocess(
+    data_directory: &'_ Path,
+    port: u32,
+) -> TmpPostgrustResult<Child> {
+    let postgres_path =
+        find_postgresql_command("bin", "postgres").expect("failed to find postgres");
+
+    Command::new(postgres_path)
+        .env("PGDATA", data_directory.to_str().unwrap())
+        .arg("-p")
+        .arg(port.to_string())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()
+        .map_err(TmpPostgrustError::SpawnSubprocessFailed)
+}
+
+#[instrument]
+pub(crate) async fn exec_init_db(data_directory: &'_ Path) -> TmpPostgrustResult<()> {
+    let initdb_path = find_postgresql_command("bin", "initdb").expect("failed to find initdb");
+
+    debug!("Initializing database in: {:?}", data_directory);
+    exec_process(
+        &mut Command::new(initdb_path)
+            .env("PGDATA", data_directory.to_str().unwrap())
+            .arg("--username=postgres"),
+        TmpPostgrustError::InitDBFailed,
+    )
+    .await
+}
+
+#[instrument]
+pub(crate) async fn exec_copy_dir(src_dir: &'_ Path, dst_dir: &'_ Path) -> TmpPostgrustResult<()> {
+    for read_dir in src_dir
+        .read_dir()
+        .map_err(TmpPostgrustError::CopyCachedInitDBFailedFileNotFound)?
+    {
+        let mut cmd = Command::new("cp");
+        #[cfg(target_os = "macos")]
+        cmd.arg("-R")
+            .arg("-c")
+            .arg(
+                read_dir
+                    .map_err(TmpPostgrustError::CopyCachedInitDBFailedFileNotFound)?
+                    .path(),
+            )
+            .arg(dst_dir);
+        #[cfg(not(target_os = "macos"))]
+        cmd.arg("-R")
+            .arg("--reflink=auto")
+            .arg(
+                read_dir
+                    .map_err(TmpPostgrustError::CopyCachedInitDBFailedFileNotFound)?
+                    .path(),
+            )
+            .arg(dst_dir);
+        exec_process(&mut cmd, TmpPostgrustError::CopyCachedInitDBFailed).await?;
+    }
+    Ok(())
+}
+
+#[instrument]
+pub(crate) async fn exec_create_db(
+    socket: &'_ Path,
+    port: u32,
+    owner: &'_ str,
+    dbname: &'_ str,
+) -> TmpPostgrustResult<()> {
+    exec_process(
+        &mut Command::new("createdb")
+            .arg("-h")
+            .arg(socket)
+            .arg("-p")
+            .arg(port.to_string())
+            .arg("-U")
+            .arg("postgres")
+            .arg("-O")
+            .arg(owner)
+            .arg("--echo")
+            .arg(dbname),
+        TmpPostgrustError::CreateDBFailed,
+    )
+    .await
+}
+
+#[instrument]
+pub(crate) async fn exec_create_user(
+    socket: &'_ Path,
+    port: u32,
+    username: &'_ str,
+) -> TmpPostgrustResult<()> {
+    exec_process(
+        &mut Command::new("createuser")
+            .arg("-h")
+            .arg(socket)
+            .arg("-p")
+            .arg(port.to_string())
+            .arg("-U")
+            .arg("postgres")
+            .arg("--superuser")
+            .arg("--echo")
+            .arg(username),
+        TmpPostgrustError::CreateDBFailed,
+    )
+    .await
+}
+
+/// ProcessGuard represents a postgresql process that is running in the background.
+/// once the guard is dropped the process will be killed.
+pub struct ProcessGuard {
+    /// Allows users to read stdout by line for debugging.
+    pub stdout_reader: Option<Lines<BufReader<ChildStdout>>>,
+    /// Allows users to read stderr by line for debugging.
+    pub stderr_reader: Option<Lines<BufReader<ChildStderr>>>,
+    /// Connection string for connecting to the temporary postgresql instance.
+    pub connection_string: String,
+
+    // Signal that the postgres process should be killed.
+    pub(crate) send_done: Option<Sender<()>>,
+    // Prevent the data directory from being dropped while
+    // the process is running.
+    pub(crate) _data_directory: TempDir,
+    // Prevent socket directory from being dropped while
+    // the process is running.
+    pub(crate) _socket_dir: Arc<TempDir>,
+    // Limit the total concurrent processes.
+    pub(crate) _process_permit: SemaphorePermit<'static>,
+}
+
+/// Signal that the process needs to end.
+impl Drop for ProcessGuard {
+    fn drop(&mut self) {
+        if let Some(sender) = self.send_done.take() {
+            sender
+                .send(())
+                .expect("failed to signal postgresql process should be killed.");
+        }
+    }
+}

src/errors.rs

@@ -0,0 +1,58 @@
+use thiserror::Error;
+
+/// UTF-8 captures of stdout and stderr for child processes used by the library.
+#[derive(Debug)]
+pub struct ProcessCapture {
+    /// Capture of stdout from the process
+    pub stdout: String,
+    /// Capture of stderr from the process
+    pub stderr: String,
+}
+
+/// Error type for possible postgresql errors.
+#[derive(Error, Debug)]
+pub enum TmpPostgrustError {
+    /// Catchall error for when a subprocess fails to run to completion
+    #[error("subprocess failed to execute")]
+    ExecSubprocessFailed {
+        /// Underlying I/O Error.
+        #[source]
+        source: std::io::Error,
+        /// Debug formatted string of Command that was attempted.
+        command: String,
+    },
+    /// Catchall error for when a subprocess fails to start
+    #[error("subprocess failed to spawn")]
+    SpawnSubprocessFailed(#[source] std::io::Error),
+    /// Error when `initdb` fails to execute.
+    #[error("initdb failed")]
+    InitDBFailed(ProcessCapture),
+    /// Error when `cp` fails for the initialized database.
+    #[error("copying cached database failed")]
+    CopyCachedInitDBFailed(ProcessCapture),
+    /// Error when a file to be copied is not found.
+    #[error("copying cached database failed, file not found")]
+    CopyCachedInitDBFailedFileNotFound(#[source] std::io::Error),
+    /// Error when a copy process cannot be joined.
+    #[cfg(feature = "tokio-process")]
+    #[error("copying cached database failed, failed to join cp process")]
+    CopyCachedInitDBFailedJoinError(#[source] tokio::task::JoinError),
+    /// Error when `createdb` fails to execute.
+    #[error("createdb failed")]
+    CreateDBFailed(ProcessCapture),
+    /// Error when `postgresql.conf` cannot be written.
+    #[error("failed to write postgresql.conf")]
+    CreateConfigFailed(#[source] std::io::Error),
+    /// Error when the PGDATA directory is empty.
+    #[error("failed to find temporary data directory")]
+    EmptyDataDirectory,
+    /// Error when the temporary unix socket directory cannot be created.
+    #[error("failed to create unix socket directory")]
+    CreateSocketDirFailed(#[source] std::io::Error),
+    /// Error when the cache directory cannot be created.
+    #[error("failed to create cache directory")]
+    CreateCacheDirFailed(#[source] std::io::Error),
+}
+
+/// Result type for `TmpPostgrustError`, used by functions in this crate.
+pub type TmpPostgrustResult<T> = Result<T, TmpPostgrustError>;