diff --git a/backend/src/main.rs b/backend/src/main.rs index 3a1899c23b..d9ae5c6e58 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -320,7 +320,7 @@ async fn windmill_main() -> anyhow::Result<()> { .unwrap_or(DEFAULT_NUM_WORKERS as i32) }; - if num_workers > 1 { + if num_workers > 1 && !std::env::var("WORKER_GROUP").is_ok_and(|x| x == "native") { println!( "We STRONGLY recommend using at most 1 worker per container, use at your own risks" ); @@ -344,8 +344,17 @@ async fn windmill_main() -> anyhow::Result<()> { }; println!("Connecting to database..."); - let db = windmill_common::connect_db(server_mode, indexer_mode).await?; + let db = windmill_common::initial_connection().await?; + let num_version = sqlx::query_scalar!("SELECT version()").fetch_one(&db).await; + + tracing::info!( + "PostgreSQL version: {} (windmill require PG >= 14)", + num_version + .ok() + .flatten() + .unwrap_or_else(|| "UNKNOWN".to_string()) + ); load_otel(&db).await; tracing::info!("Database connected"); @@ -362,16 +371,6 @@ async fn windmill_main() -> anyhow::Result<()> { let _guard = windmill_common::tracing_init::initialize_tracing(&hostname, &mode, &environment); - let num_version = sqlx::query_scalar!("SELECT version()").fetch_one(&db).await; - - tracing::info!( - "PostgreSQL version: {} (windmill require PG >= 14)", - num_version - .ok() - .flatten() - .unwrap_or_else(|| "UNKNOWN".to_string()) - ); - let is_agent = mode == Mode::Agent; #[cfg(feature = "parquet")] @@ -379,7 +378,7 @@ async fn windmill_main() -> anyhow::Result<()> { .ok() .is_some_and(|x| x == "1" || x == "true"); - if !is_agent { + if !is_agent && !indexer_mode { let skip_migration = std::env::var("SKIP_MIGRATION") .map(|val| val == "true") .unwrap_or(false); @@ -392,6 +391,11 @@ async fn windmill_main() -> anyhow::Result<()> { } } + drop(db); + let worker_mode = num_workers > 0; + + let db = windmill_common::connect_db(server_mode, indexer_mode, worker_mode).await?; + let (killpill_tx, mut killpill_rx) = tokio::sync::broadcast::channel::<()>(2); let mut monitor_killpill_rx = killpill_tx.subscribe(); let (killpill_phase2_tx, _killpill_phase2_rx) = tokio::sync::broadcast::channel::<()>(2); @@ -455,8 +459,6 @@ Windmill Community Edition {GIT_VERSION} } } - let worker_mode = num_workers > 0; - if server_mode || worker_mode || indexer_mode { let port_var = std::env::var("PORT").ok().and_then(|x| x.parse().ok()); diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 550745d31a..4bf68f59f9 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -221,28 +221,42 @@ async fn reset() -> () { todo!() } -pub async fn connect_db( - server_mode: bool, - indexer_mode: bool, -) -> anyhow::Result> { - use anyhow::Context; +pub async fn get_database_url() -> Result { use std::env::var; use tokio::fs::File; use tokio::io::AsyncReadExt; - - let database_url = match var("DATABASE_URL_FILE") { + match var("DATABASE_URL_FILE") { Ok(file_path) => { let mut file = File::open(file_path).await?; let mut contents = String::new(); file.read_to_string(&mut contents).await?; - contents.trim().to_string() + Ok(contents.trim().to_string()) } Err(_) => var("DATABASE_URL").map_err(|_| { Error::BadConfig( "Either DATABASE_URL_FILE or DATABASE_URL env var is missing".to_string(), ) - })?, - }; + }), + } +} + +pub async fn initial_connection() -> Result, error::Error> { + let database_url = get_database_url().await?; + sqlx::postgres::PgPoolOptions::new() + .max_connections(2) + .connect_with(sqlx::postgres::PgConnectOptions::from_str(&database_url)?) + .await + .map_err(|err| Error::ConnectingToDatabase(err.to_string())) +} + +pub async fn connect_db( + server_mode: bool, + indexer_mode: bool, + worker_mode: bool, +) -> anyhow::Result> { + use anyhow::Context; + + let database_url = get_database_url().await?; let max_connections = match std::env::var("DATABASE_CONNECTIONS") { Ok(n) => n.parse::().context("invalid DATABASE_CONNECTIONS")?, @@ -263,12 +277,13 @@ pub async fn connect_db( } }; - Ok(connect(&database_url, max_connections).await?) + Ok(connect(&database_url, max_connections, worker_mode).await?) } pub async fn connect( database_url: &str, max_connections: u32, + worker_mode: bool, ) -> Result, error::Error> { use std::time::Duration; @@ -276,6 +291,18 @@ pub async fn connect( .min_connections((max_connections / 5).clamp(3, max_connections)) .max_connections(max_connections) .max_lifetime(Duration::from_secs(30 * 60)) // 30 mins + .after_connect(move |conn, _| { + if worker_mode { + Box::pin(async move { + sqlx::query("SET enable_seqscan = OFF;") + .execute(conn) + .await?; + Ok(()) + }) + } else { + Box::pin(async move { Ok(()) }) + } + }) .connect_with( sqlx::postgres::PgConnectOptions::from_str(database_url)?.statement_cache_capacity(400), ) diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index 4478601e7e..23b5ed8f3e 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -104,7 +104,7 @@ lazy_static::lazy_static! { } fn format_pull_query(peek: String) -> String { - format!( + let r = format!( "WITH peek AS ( {} ), q AS NOT MATERIALIZED ( @@ -145,7 +145,9 @@ fn format_pull_query(peek: String) -> String { FROM q, j LEFT JOIN v2_job_status f USING (id)", peek - ) + ); + tracing::debug!("pull query: {}", r); + r } pub async fn make_suspended_pull_query(wc: &WorkerConfig) {