Rework pg-embed feature.

gz · ryzhyk · commit 762d382e4c78 · 2023-05-02T09:36:24.000-07:00
Turns out if it's in a once-cell drop doesn't get called if the
developer hits ctrl+c which means next time the manager is restarted
it won't start again. So we put it back in ProjectDb.
Also removed oncecell dependency which means we can simplify feature
to just 'pg-embed'.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/pipeline_manager/Cargo.toml b/crates/pipeline_manager/Cargo.toml
@@ -37,7 +37,6 @@ tokio-postgres = "0.7"
 async-trait = "0.1"
 # Waiting for https://github.com/faokunega/pg-embed/pull/26
 pg-embed = { git = "https://github.com/gz/pg-embed.git", rev = "8906af8", optional = true }
-once_cell = { version = "1.17.1", optional = true }
 
 [target.'cfg(unix)'.dependencies]
 daemonize = { version = "0.4.1" }
@@ -50,10 +49,7 @@ proptest = "1.0.0"
 proptest-derive = "0.3.0"
 pretty_assertions = "1.3.0"
 # Workaround to enable dev feature during tests: https://github.com/rust-lang/cargo/issues/2911
-dbsp_pipeline_manager = { path = ".", features = ["dev"]}
+dbsp_pipeline_manager = { path = ".", features = ["pg-embed"]}
 
 [package.metadata.cargo-udeps.ignore]
-development = ["dbsp_pipeline_manager"] # false positive from cargo udeps
-
-[features]
-dev = ["dep:pg-embed", "dep:once_cell"]
+development = ["dbsp_pipeline_manager"] # false positive from cargo udeps
diff --git a/crates/pipeline_manager/src/config.rs b/crates/pipeline_manager/src/config.rs
@@ -228,7 +228,7 @@ impl ManagerConfig {
     }
 
     /// Where Postgres embed stores the database.
-    #[cfg(feature = "dev")]
+    #[cfg(feature = "pg-embed")]
     pub(crate) fn postgres_embed_data_dir(&self) -> PathBuf {
         Path::new(&self.working_directory).join("data")
     }
diff --git a/crates/pipeline_manager/src/db/mod.rs b/crates/pipeline_manager/src/db/mod.rs
@@ -12,7 +12,7 @@ use utoipa::ToSchema;
 #[cfg(test)]
 mod test;
 
-#[cfg(any(test, feature = "dev"))]
+#[cfg(any(test, feature = "pg-embed"))]
 mod pg_setup;
 pub(crate) mod storage;
 
@@ -29,13 +29,13 @@ pub(crate) mod storage;
 /// time, which determines the position of the project in the queue.
 pub(crate) struct ProjectDB {
     conn: Client,
+    // Used in dev mode for having an embedded Postgres DB live through the
+    // lifetime of the program.
+    #[cfg(feature = "pg-embed")]
+    #[allow(dead_code)] // It has to stay alive until ProjectDB is dropped.
+    pg_inst: Option<pg_embed::postgres::PgEmbed>,
 }
 
-// Used in dev mode for having an embedded Postgres DB live through the lifetime
-// of the program.
-#[cfg(feature = "dev")]
-static PG: once_cell::sync::OnceCell<pg_embed::postgres::PgEmbed> =
-    once_cell::sync::OnceCell::new();
 
 /// Unique project id.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Ord, PartialOrd, Serialize, Deserialize, ToSchema)]
@@ -1021,18 +1021,16 @@ impl ProjectDB {
         let connection_str = config.database_connection_string();
         let initial_sql = &config.initial_sql;
 
-        #[cfg(feature = "dev")]
-        let connection_str = if connection_str.starts_with("postgres-embed") {
+        #[cfg(feature = "pg-embed")]
+        if connection_str.starts_with("postgres-embed") {
             let database_dir = config.postgres_embed_data_dir();
-            let pg = pg_setup::install(database_dir, true, Some(8082)).await?;
-            let connection_string = pg.db_uri.to_string();
-            let _ = PG.set(pg);
-            connection_string
-        } else {
-            connection_str
+            let pg_inst = pg_setup::install(database_dir, true, Some(8082)).await?;
+            let connection_string = pg_inst.db_uri.to_string();
+            return Self::connect_inner(connection_string.as_str(), initial_sql, Some(pg_inst)).await
         };
 
-        Self::connect_inner(connection_str.as_str(), initial_sql).await
+        Self::connect_inner(connection_str.as_str(), initial_sql, #[cfg(feature = "pg-embed")] None).await
+
     }
 
     /// Connect to the project database.
@@ -1046,7 +1044,7 @@ impl ProjectDB {
     /// - `is_persistent`: Whether the embedded postgres database should be
     ///   persistent or removed on shutdown.
     /// - `port`: The port to use for the embedded Postgres database to run on.
-    async fn connect_inner(connection_str: &str, initial_sql: &Option<String>) -> AnyResult<Self> {
+    async fn connect_inner(connection_str: &str, initial_sql: &Option<String>, #[cfg(feature = "pg-embed")] pg_inst: Option<pg_embed::postgres::PgEmbed>) -> AnyResult<Self> {
         if !connection_str.starts_with("postgres") {
             panic!("Unsupported connection string {}", connection_str)
         }
@@ -1162,7 +1160,10 @@ impl ProjectDB {
             }
         }
 
-        Ok(Self { conn: client })
+        #[cfg(feature = "pg-embed")]
+        return Ok(Self { conn: client, pg_inst });
+        #[cfg(not(feature = "pg-embed"))]
+        return Ok(Self { conn: client });
     }
 
     /// Attach connector to the config.
diff --git a/crates/pipeline_manager/src/db/test.rs b/crates/pipeline_manager/src/db/test.rs
@@ -7,7 +7,6 @@ use crate::db::{pg_setup, DBError};
 use anyhow::Result as AnyResult;
 use async_trait::async_trait;
 use chrono::DateTime;
-use pg_embed::postgres::PgEmbed;
 use pretty_assertions::assert_eq;
 use proptest::prelude::*;
 use proptest::test_runner::{Config, TestRunner};
@@ -21,7 +20,6 @@ use tokio::sync::Mutex;
 
 struct DbHandle {
     db: ProjectDB,
-    pg: PgEmbed,
     _temp_dir: TempDir,
 }
 
@@ -31,9 +29,11 @@ impl Drop for DbHandle {
         // shutdown postgres). Otherwise postgres log an error that the
         // directory is already gone during shutdown which could be
         // confusing for a developer.
-        let _r = async {
-            self.pg.stop_db().await.unwrap();
-        };
+        if let Some(pg) = self.db.pg_inst.as_mut() {
+            let _r = async {
+                pg.stop_db().await.unwrap();
+            };
+        }
     }
 }
 
@@ -58,14 +58,13 @@ async fn test_setup() -> DbHandle {
     let pg = pg_setup::install(temp_path.into(), false, Some(port))
         .await
         .unwrap();
-
-    let conn = ProjectDB::connect_inner(&pg.db_uri, &Some("".to_string()))
+    let db_uri = pg.db_uri.clone();
+    let conn = ProjectDB::connect_inner(&db_uri, &Some("".to_string()), Some(pg))
         .await
         .unwrap();
 
     DbHandle {
         db: conn,
-        pg: pg,
         _temp_dir,
     }
 }
diff --git a/scripts/start_manager.sh b/scripts/start_manager.sh
@@ -30,8 +30,8 @@ pkill -9 dbsp_pipeline_
 
 set -e
 
-cd "${MANAGER_DIR}" && ~/.cargo/bin/cargo build --release
-cd "${MANAGER_DIR}" && ~/.cargo/bin/cargo run --release -- \
+cd "${MANAGER_DIR}" && ~/.cargo/bin/cargo build --release --features pg-embed
+cd "${MANAGER_DIR}" && ~/.cargo/bin/cargo run --release --features pg-embed -- \
     --bind-address="${BIND_ADDRESS}" \
     --working-directory="${WORKING_DIR_ABS}" \
     --sql-compiler-home="${SQL_COMPILER_DIR}" \

Original file line number	Diff line number	Diff line change
`@@ -228,7 +228,7 @@ impl ManagerConfig {`
`228`	`228`	`}`
`229`	`229`
`230`	`230`	`/// Where Postgres embed stores the database.`
`231`		`- #[cfg(feature = "dev")]`
	`231`	`+ #[cfg(feature = "pg-embed")]`
`232`	`232`	`pub(crate) fn postgres_embed_data_dir(&self) -> PathBuf {`
`233`	`233`	`Path::new(&self.working_directory).join("data")`
`234`	`234`	`}`