diff --git a/.github/workflows/build-test.yaml b/.github/workflows/build-test.yaml index 1e96e81..ec11166 100644 --- a/.github/workflows/build-test.yaml +++ b/.github/workflows/build-test.yaml @@ -16,8 +16,22 @@ jobs: run: cargo fmt --all -- --check - name: Run cargo clippy run: make lint - test: - name: Unit Tests + + test-windows: + name: Windows Unit Tests + runs-on: windows-latest + steps: + - name: Install postgres + uses: tj-actions/install-postgresql@v3 + with: + postgresql-version: 17 + - name: Checkout + uses: actions/checkout@v4 + - name: Run tests + run: make test-ci-windows + + test-unix: + name: Unix Unit Tests runs-on: ubuntu-latest steps: - name: Install postgres @@ -27,7 +41,7 @@ jobs: - name: Checkout uses: actions/checkout@v4 - name: Run tests - run: make test-ci + run: make test-ci-unix doc: name: Documentation Check diff --git a/Cargo.lock b/Cargo.lock index ad22823..f217ebf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "addr2line" diff --git a/Makefile b/Makefile index f60a8a9..8682b2c 100644 --- a/Makefile +++ b/Makefile @@ -9,11 +9,15 @@ test: cargo test --tests --examples --all-features # Run all tests with and without all features -test-ci: +test-ci-unix: cargo test --tests --examples --no-default-features cargo test --tests --examples --all-features make -C examples/python-sqlalchemy test-ci +test-ci-windows: + cargo test --tests --examples --no-default-features + cargo test --tests --examples --all-features + # Run clippy lint: cargo clippy --no-deps --all-targets --all-features -- -W clippy::pedantic \ diff --git a/src/daemon.rs b/src/daemon.rs index aa915a4..b1f9903 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -2,9 +2,8 @@ use crate::{PgTempDB, PgTempDBBuilder}; use std::net::SocketAddr; -use tokio::io::AsyncWriteExt; +use tokio::io::{self, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use tokio::signal::unix::{signal, SignalKind}; #[cfg(feature = "cli")] /// Contains the clap args struct @@ -144,7 +143,7 @@ impl PgTempDaemon { /// Start the daemon, listening for either TCP connections on the configured port. The server /// shuts down when sent a SIGINT (e.g. via ctrl-C). - pub async fn start(mut self) { + pub async fn start(self) { let uri = self.conn_uri(); if self.single_mode { println!("starting pgtemp server in single mode at {}", uri); @@ -155,39 +154,34 @@ impl PgTempDaemon { let listener = TcpListener::bind(("127.0.0.1", self.port)) .await .expect("failed to bind to daemon port"); + self.listen(listener).await + } + + /// Main daemon listening loop for unix + #[cfg(unix)] + async fn listen(mut self, listener: TcpListener) { + use tokio::signal::unix::{signal, SignalKind}; + let mut sig = signal(SignalKind::interrupt()).expect("failed to hook to interrupt signal"); loop { tokio::select! { - res = listener.accept() => { - if let Ok((client_conn, client_addr)) = res { - client_conn.set_nodelay(true).expect("failed to set nodelay on client connection"); - let db: Option; - let db_port: u16; - if self.single_mode { - db = None; - db_port = self.dbs[0].db_port(); - } - else { - let take_db = self.dbs.pop().unwrap(); - db_port = take_db.db_port(); - db = Some(take_db); - } - let db_conn = TcpStream::connect(("127.0.0.1", db_port)) - .await - .expect("failed to connect to postgres server"); - db_conn - .set_nodelay(true) - .expect("failed to set nodelay on db connection"); - tokio::spawn(async move { proxy_connection(db, db_conn, client_conn, client_addr).await }); - // preallocate a new db after one is used - if self.dbs.is_empty() && !self.single_mode { - self.allocate_db().await; - } - } - else { - println!("idk when this errs"); - } + res = listener.accept() => self.on_listener_accept(res).await, + _sig_event = sig.recv() => { + println!("got interrupt, exiting"); + break; } + } + } + } + + /// Main daemon listening loop for windows + #[cfg(windows)] + async fn listen(mut self, listener: TcpListener) { + let mut sig = + tokio::signal::windows::ctrl_c().expect("failed to hook windows interrupt signal"); + loop { + tokio::select! { + res = listener.accept() => self.on_listener_accept(res).await, _sig_event = sig.recv() => { println!("got interrupt, exiting"); break; @@ -195,6 +189,40 @@ impl PgTempDaemon { } } } + + /// Called when a connection is accepted from a TcpListener. + async fn on_listener_accept(&mut self, result: io::Result<(TcpStream, SocketAddr)>) { + if let Ok((client_conn, client_addr)) = result { + client_conn + .set_nodelay(true) + .expect("failed to set nodelay on client connection"); + let db: Option; + let db_port: u16; + if self.single_mode { + db = None; + db_port = self.dbs[0].db_port(); + } else { + let take_db = self.dbs.pop().unwrap(); + db_port = take_db.db_port(); + db = Some(take_db); + } + let db_conn = TcpStream::connect(("127.0.0.1", db_port)) + .await + .expect("failed to connect to postgres server"); + db_conn + .set_nodelay(true) + .expect("failed to set nodelay on db connection"); + tokio::spawn( + async move { proxy_connection(db, db_conn, client_conn, client_addr).await }, + ); + // preallocate a new db after one is used + if self.dbs.is_empty() && !self.single_mode { + self.allocate_db().await; + } + } else { + println!("idk when this errs"); + } + } } /// When we're in single mode, we pass None to the db here so it doesn't get deallocated when the diff --git a/src/lib.rs b/src/lib.rs index 0db4c38..1c8f165 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -161,7 +161,6 @@ impl PgTempDB { .postgres_process .take() .expect("shutdown with no postgres process"); - let temp_dir = self.temp_dir.take().unwrap(); // fast (not graceful) shutdown via SIGINT // TODO: graceful shutdown via SIGTERM @@ -173,11 +172,28 @@ impl PgTempDB { // the postgres server says "we're still connected to a client, can't shut down yet" and we // have a deadlock. #[allow(clippy::cast_possible_wrap)] - let _ret = unsafe { libc::kill(postgres_process.id() as i32, libc::SIGINT) }; + #[cfg(unix)] + { + unsafe { + libc::kill(postgres_process.id() as i32, libc::SIGINT); + } + } + #[cfg(windows)] + { + std::process::Command::new("pg_ctl") + .arg("stop") + .arg("-D") + .arg(self.data_dir()) + .output() + .expect("Failed to stop server with pg_ctl. Is it installed and on your path?"); + } + let _output = postgres_process .wait_with_output() .expect("postgres server failed to exit cleanly"); + let temp_dir = self.temp_dir.take().unwrap(); + if self.persist { // this prevents the dir from being deleted on drop let _path = temp_dir.into_path(); diff --git a/src/run_db.rs b/src/run_db.rs index f269162..2a98eb1 100644 --- a/src/run_db.rs +++ b/src/run_db.rs @@ -1,4 +1,7 @@ +use std::collections::HashMap; +use std::ffi::OsStr; use std::{ + path::PathBuf, process::{Child, Command}, time::Duration, }; @@ -9,10 +12,34 @@ use crate::PgTempDBBuilder; const CREATEDB_MAX_TRIES: u32 = 10; const CREATEDB_RETRY_DELAY: Duration = Duration::from_millis(100); +#[cfg(unix)] fn current_user_is_root() -> bool { unsafe { libc::getuid() == 0 } } +#[cfg(unix)] +fn get_command(program: S) -> Command +where + S: AsRef, +{ + // postgres will not run as root, so try to run initdb as postgres user if we are root so that + // when running the server as the postgres user it can access the files + if current_user_is_root() { + let mut cmd = Command::new("sudo"); + cmd.args(["-u", "postgres"]).arg(program); + cmd + } else { + Command::new(program) + } +} +#[cfg(windows)] +fn get_command(program: S) -> Command +where + S: AsRef, +{ + Command::new(program) +} + /// Execute the `initdb` binary with the parameters configured in PgTempDBBuilder. pub fn init_db(builder: &mut PgTempDBBuilder) -> TempDir { let temp_dir = { @@ -23,6 +50,7 @@ pub fn init_db(builder: &mut PgTempDBBuilder) -> TempDir { } }; + #[cfg(unix)] // if current user is root, data dir etc need to be owned by postgres user if current_user_is_root() { // TODO: don't shell out to chown, get the userid of postgres and just call std::os @@ -59,13 +87,7 @@ pub fn init_db(builder: &mut PgTempDBBuilder) -> TempDir { // postgres will not run as root, so try to run initdb as postgres user if we are root so that // when running the server as the postgres user it can access the files - let mut cmd: Command; - if current_user_is_root() { - cmd = Command::new("sudo"); - cmd.args(["-u", "postgres"]).arg(initdb_path); - } else { - cmd = Command::new(initdb_path); - } + let mut cmd: Command = get_command(initdb_path); cmd.args(["-D", data_dir_str]) .arg("-N") // no fsync, starts slightly faster @@ -90,24 +112,45 @@ pub fn init_db(builder: &mut PgTempDBBuilder) -> TempDir { temp_dir } -pub fn run_db(temp_dir: &TempDir, mut builder: PgTempDBBuilder) -> Child { - let data_dir = temp_dir.path().join("pg_data_dir"); - let data_dir_str = data_dir.to_str().unwrap(); - let port = builder.get_port_or_set_random(); +#[cfg(windows)] +/// waits for a db to be ready +/// if the server never started, returns the last command output +fn wait_for_db_ready( + bin_path: Option<&PathBuf>, + port: u16, + max_retries: u32, + retry_delay: Duration, +) -> Option { + let mut isready_last_error_output = None; + let isready_path = bin_path.map_or("pg_isready".into(), |p| p.join("pg_isready")); - // postgres will not run as root, so try to run as postgres if we are root - let postgres_path = builder - .bin_path - .as_ref() - .map_or("postgres".into(), |p| p.join("postgres")); - let mut pgcmd: Command; - if current_user_is_root() { - pgcmd = Command::new("sudo"); - pgcmd.args(["-u", "postgres"]).arg(postgres_path); - } else { - pgcmd = Command::new(postgres_path); - }; + let mut server_status_cmd = Command::new(&isready_path); + server_status_cmd.args(["-p", &port.to_string()]); + + for _ in 0..max_retries { + let server_status = server_status_cmd.output().unwrap(); + + if server_status.status.success() { + isready_last_error_output = None; + break; + } + isready_last_error_output = Some(server_status); + std::thread::sleep(retry_delay); + } + + isready_last_error_output +} + +#[cfg(unix)] +fn get_run_db_cmd( + bin_path: Option<&PathBuf>, + data_dir_str: &str, + port: u16, + server_configs: &HashMap, +) -> Command { + let postgres_path = bin_path.map_or("postgres".into(), |p| p.join("postgres")); + let mut pgcmd = get_command(postgres_path); pgcmd .args(["-c", &format!("unix_socket_directories={}", data_dir_str)]) .args(["-c", &format!("port={port}")]) @@ -117,10 +160,54 @@ pub fn run_db(temp_dir: &TempDir, mut builder: PgTempDBBuilder) -> Child { .args(["-c", "synchronous_commit=off"]) .args(["-c", "full_page_writes=off"]) .args(["-c", "autovacuum=off"]) - .args(["-D", data_dir.to_str().unwrap()]); - for (key, val) in &builder.server_configs { + .args(["-D", data_dir_str]); + for (key, val) in server_configs { pgcmd.args(["-c", &format!("{}={}", key, val)]); } + pgcmd +} +#[cfg(windows)] +fn get_run_db_cmd( + bin_path: Option<&PathBuf>, + data_dir_str: &str, + port: u16, + server_configs: &HashMap, +) -> Command { + let pg_ctl_path = bin_path.map_or("pg_ctl".into(), |p| p.join("pg_ctl")); + + // build postgres command args + let mut pg_cmd_args = vec![ + format!("-c unix_socket_directories={}", data_dir_str), + format!("-c port={}", port), + "-c fsync=off".into(), + "-c synchronous_commit=off".into(), + "-c full_page_writes=off".into(), + "-c autovacuum=off".into(), + ]; + for (key, val) in server_configs { + pg_cmd_args.push(format!("-c {}={}", key, val)); + } + + let mut pgcmd = get_command(pg_ctl_path); + pgcmd + .arg("start") + .args(["-D", data_dir_str]) + .args(["-o", &pg_cmd_args.join(" ")]); + pgcmd +} + +pub fn run_db(temp_dir: &TempDir, mut builder: PgTempDBBuilder) -> Child { + let data_dir = temp_dir.path().join("pg_data_dir"); + let data_dir_str = data_dir.to_str().unwrap(); + let port = builder.get_port_or_set_random(); + + // postgres will not run as root, so try to run as postgres if we are root + let mut pgcmd: Command = get_run_db_cmd( + builder.bin_path.as_ref(), + data_dir_str, + port, + &builder.server_configs, + ); // don't output postgres output to stdout/stderr pgcmd @@ -131,6 +218,26 @@ pub fn run_db(temp_dir: &TempDir, mut builder: PgTempDBBuilder) -> Child { .spawn() .expect("Failed to start postgres. Is it installed and on your path?"); + // wait for db to be started + #[cfg(windows)] + { + let db_ready_error = wait_for_db_ready( + builder.bin_path.as_ref(), + port, + CREATEDB_MAX_TRIES, + CREATEDB_RETRY_DELAY, + ); + if let Some(output) = db_ready_error { + let stdout = output.stdout; + let stderr = output.stderr; + panic!( + "db did not start! stdout: {}\n\nstderr: {}", + String::from_utf8_lossy(&stdout), + String::from_utf8_lossy(&stderr) + ); + } + } + #[cfg(unix)] std::thread::sleep(CREATEDB_RETRY_DELAY); let user = builder.get_user(); diff --git a/tests/startup.rs b/tests/startup.rs index 3f6f1ca..a3a1438 100644 --- a/tests/startup.rs +++ b/tests/startup.rs @@ -1,7 +1,12 @@ //! Basic startup/shutdown tests -use pgtemp::{PgTempDB, PgTempDBBuilder}; +use pgtemp::PgTempDB; + +#[cfg(unix)] +use pgtemp::PgTempDBBuilder; +#[cfg(unix)] use std::{io::Write, os::unix::fs::OpenOptionsExt}; +#[cfg(unix)] use tempfile::TempDir; #[test] @@ -67,6 +72,7 @@ fn test_tempdb_bringup_shutdown_persist() { #[test] #[should_panic(expected = "this is not initdb")] +#[cfg(unix)] /// Start a database by specifying the bin_path fn test_tempdb_bin_path() { use std::io::Write; @@ -88,6 +94,7 @@ fn test_tempdb_bin_path() { } #[test] +#[cfg(unix)] fn test_slow_postgres_startup() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let dir_path = temp_dir.path().to_owned();