From 1501dbd5a5f7d25d0c5cecde496dff763edde529 Mon Sep 17 00:00:00 2001 From: Bojan Serafimov Date: Thu, 18 Aug 2022 15:16:14 -0400 Subject: [PATCH] Cleanup --- Cargo.lock | 1 + integration_tests/src/basic.rs | 19 +++-- libs/pg_bin/Cargo.toml | 1 + libs/pg_bin/src/lib.rs | 113 +++++++++++++++++---------- libs/utils/src/command_extensions.rs | 21 +++++ libs/utils/src/lib.rs | 3 + libs/utils/src/output_capture.rs | 10 +++ libs/utils/src/pg_bin.rs | 0 8 files changed, 116 insertions(+), 52 deletions(-) create mode 100644 libs/utils/src/command_extensions.rs create mode 100644 libs/utils/src/output_capture.rs delete mode 100644 libs/utils/src/pg_bin.rs diff --git a/Cargo.lock b/Cargo.lock index f6e9714a18..e680f94612 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1994,6 +1994,7 @@ name = "pg_bin" version = "0.1.0" dependencies = [ "tokio-postgres", + "utils", ] [[package]] diff --git a/integration_tests/src/basic.rs b/integration_tests/src/basic.rs index 0702f52f3b..699ac86ff7 100644 --- a/integration_tests/src/basic.rs +++ b/integration_tests/src/basic.rs @@ -1,26 +1,25 @@ #[cfg(test)] mod tests { - use pg_bin::LocalPostgres; + use pg_bin::PgDatadir; use std::path::PathBuf; use tokio_postgres::NoTls; #[tokio::test] async fn test_postgres_select_1() -> anyhow::Result<()> { // Test setup - let pg_datadir = PathBuf::from("/home/bojan/tmp/t1/"); + let output = PathBuf::from("/home/bojan/tmp/"); let pg_prefix = PathBuf::from("/home/bojan/src/neondatabase/neon/tmp_install/bin/"); - // Get a postgres - let mut postgres = LocalPostgres::new(pg_datadir, pg_prefix); - postgres.start(); - let config = postgres.admin_conn_info(); + // Init datadir + let pg_datadir_path = PathBuf::from("/home/bojan/tmp/t1/"); + let pg_datadir = PgDatadir::new_initdb(pg_datadir_path, &pg_prefix, &output, true); - if let Err(e) = config.connect(NoTls).await { - eprintln!("error error {:?}", e); - } + // Get a postgres + let postgres = pg_datadir.spawn_postgres(pg_prefix, output); + let conn_info = postgres.admin_conn_info(); // Get client, run connection - let (client, connection) = config.connect(NoTls).await?; + let (client, connection) = conn_info.connect(NoTls).await?; tokio::spawn(async move { if let Err(e) = connection.await { eprintln!("connection error: {}", e); diff --git a/libs/pg_bin/Cargo.toml b/libs/pg_bin/Cargo.toml index d26ae238ba..d614f866c1 100644 --- a/libs/pg_bin/Cargo.toml +++ b/libs/pg_bin/Cargo.toml @@ -7,3 +7,4 @@ edition = "2021" [dependencies] tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="d052ee8b86fff9897c77b0fe89ea9daba0e1fa38" } +utils = { path = "../utils" } diff --git a/libs/pg_bin/src/lib.rs b/libs/pg_bin/src/lib.rs index 3febee7f3b..bdb0fb5f04 100644 --- a/libs/pg_bin/src/lib.rs +++ b/libs/pg_bin/src/lib.rs @@ -1,57 +1,81 @@ //! Utils for runnig postgres binaries as subprocesses. -use std::{fs::{File, remove_dir_all}, path::PathBuf, process::{Child, Command, Stdio}, time::Duration}; +use std::{fs::{File, remove_dir_all}, path::PathBuf, process::{Child, Command}, time::Duration}; use std::io::Write; +use utils::command_extensions::NeonCommandExtensions; -pub struct LocalPostgres { - datadir: PathBuf, - pg_prefix: PathBuf, - port: u16, - running: Option, + +pub struct PgDatadir { + path: PathBuf } -impl LocalPostgres { - pub fn new(datadir: PathBuf, pg_prefix: PathBuf) -> Self { - LocalPostgres { - datadir, - pg_prefix, - port: 54729, - running: None, +impl PgDatadir { + pub fn new_initdb( + path: PathBuf, + pg_prefix: &PathBuf, + command_output_dir: &PathBuf, + remove_if_exists: bool + ) -> Self { + if remove_if_exists { + remove_dir_all(path.clone()).ok(); } - } - // TODO is this the right interface? Why not start it? - pub fn start(&mut self) { - remove_dir_all(self.datadir.as_os_str()).ok(); - - let status = Command::new(self.pg_prefix.join("initdb")) + let status = Command::new(pg_prefix.join("initdb")) .arg("-D") - .arg(self.datadir.as_os_str()) - .stdout(Stdio::null()) // TODO to file instead - .stderr(Stdio::null()) // TODO to file instead + .arg(path.clone()) + .capture_to_files(command_output_dir.clone(), "initdb") .status() .expect("failed to get status"); assert!(status.success()); - // Write conf - let mut conf = File::create(self.datadir.join("postgresql.conf")) - .expect("failed to create file"); - writeln!(&mut conf, "port = {}", self.port) - .expect("failed to write conf"); - - // TODO check if already running - let out_file = File::create(self.datadir.join("pg.log")).expect("can't make file"); - let err_file = File::create(self.datadir.join("pg.err")).expect("can't make file"); - self.running.replace(Command::new(self.pg_prefix.join("postgres")) - .env("PGDATA", self.datadir.as_os_str()) - .stdout(Stdio::from(out_file)) - .stderr(Stdio::from(err_file)) - .spawn() - .expect("postgres failed to spawn")); - - std::thread::sleep(Duration::from_millis(300)); + Self { + path + } } + pub fn load_existing(path: PathBuf) -> Self{ + Self { + path + } + } + + pub fn path(&self) -> PathBuf { + self.path.clone() + } + + pub fn spawn_postgres(self, pg_prefix: PathBuf, command_output_dir: PathBuf) -> LocalPostgres { + let port = 54729; + + // Write conf + // TODO don't override existing conf + // - instead infer port from conf + let mut conf = File::create(self.path().join("postgresql.conf")).expect("failed to create file"); + writeln!(&mut conf, "port = {}", port).expect("failed to write conf"); + + let process = Command::new(pg_prefix.join("postgres")) + .env("PGDATA", self.path()) + .capture_to_files(command_output_dir, "pg") + .spawn() + .expect("postgres failed to spawn"); + + // Wait until ready. TODO improve this + std::thread::sleep(Duration::from_millis(300)); + + LocalPostgres { + datadir: self, + port: 54729, + process, + } + } +} + +pub struct LocalPostgres { + datadir: PgDatadir, + port: u16, + process: Child, +} + +impl LocalPostgres { pub fn admin_conn_info(&self) -> tokio_postgres::Config { // I don't like this, but idk what else to do let whoami = Command::new("whoami").output().unwrap().stdout; @@ -66,12 +90,17 @@ impl LocalPostgres { .user(&user); config } + + pub fn stop(mut self) -> PgDatadir { + self.process.kill().expect("failed to kill child"); + PgDatadir { + path: self.datadir.path.clone() + } + } } impl Drop for LocalPostgres { fn drop(&mut self) { - if let Some(mut child) = self.running.take() { - child.kill().expect("failed to kill child"); - } + self.process.kill().expect("failed to kill child"); } } diff --git a/libs/utils/src/command_extensions.rs b/libs/utils/src/command_extensions.rs new file mode 100644 index 0000000000..ef7dcae3fe --- /dev/null +++ b/libs/utils/src/command_extensions.rs @@ -0,0 +1,21 @@ +use std::path::PathBuf; +use std::{os::unix::prelude::CommandExt, process::Command}; +use std::fs::File; + + +pub trait NeonCommandExtensions: CommandExt { + fn capture_to_files(&mut self, path: PathBuf, name: &str) -> &mut Command; +} + +impl NeonCommandExtensions for Command { + fn capture_to_files(&mut self, path: PathBuf, name: &str) -> &mut Command { + let out_file = File::create(path.join(format!("{}.out", name))) + .expect("can't make file"); + let err_file = File::create(path.join(format!("{}.out", name))) + .expect("can't make file"); + + // TODO touch files? + + self.stdout(out_file).stderr(err_file) + } +} diff --git a/libs/utils/src/lib.rs b/libs/utils/src/lib.rs index 1b011bb73a..9e447b9537 100644 --- a/libs/utils/src/lib.rs +++ b/libs/utils/src/lib.rs @@ -54,6 +54,9 @@ pub mod nonblock; // Default signal handling pub mod signals; +// Helpers for running commands +pub mod command_extensions; + /// This is a shortcut to embed git sha into binaries and avoid copying the same build script to all packages /// /// we have several cases: diff --git a/libs/utils/src/output_capture.rs b/libs/utils/src/output_capture.rs new file mode 100644 index 0000000000..67d6e4bf88 --- /dev/null +++ b/libs/utils/src/output_capture.rs @@ -0,0 +1,10 @@ +use std::{os::unix::prelude::CommandExt, process::Command}; + + + +/// Command with ability to capture stdout and stderr to files +trait CaptureOutputToFile: CommandExt { + fn capture_to_file(&mut self) -> &mut Command; +} + +impl CaptureOutputToFile diff --git a/libs/utils/src/pg_bin.rs b/libs/utils/src/pg_bin.rs deleted file mode 100644 index e69de29bb2..0000000000