diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 5c08ebe06a..acc8607dce 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -630,25 +630,27 @@ impl ComputeNode { /// Start Postgres as a child process and manage DBs/roles. /// After that this will hang waiting on the postmaster process to exit. #[instrument(skip_all)] - pub fn start_postgres( - &self, - storage_auth_token: Option, - ) -> Result { + pub fn start_postgres(&self, pspec: &ParsedSpec) -> Result { let pgdata_path = Path::new(&self.pgdata); + let spec = &pspec.spec; // Run postgres as a child process. - let mut pg = maybe_cgexec(&self.pgbin) - .args(["-D", &self.pgdata]) - .envs(if let Some(storage_auth_token) = &storage_auth_token { - vec![("NEON_AUTH_TOKEN", storage_auth_token)] - } else { - vec![] - }) - .spawn() - .expect("cannot start postgres process"); + let mut pg_command = maybe_cgexec(&self.pgbin); + pg_command.args(["-D", &self.pgdata]); + if let Some(storage_auth_token) = &spec.storage_auth_token { + pg_command.env("NEON_AUTH_TOKEN", storage_auth_token); + } + if let Some(env_vars) = &spec.env_vars { + pg_command.envs( + env_vars + .iter() + .map(|(k, v)| (k, v.clone().unwrap_or(String::new()))), + ); + } + + let mut pg = pg_command.spawn().expect("cannot start postgres process"); wait_for_postgres(&mut pg, pgdata_path)?; - Ok(pg) } @@ -797,7 +799,7 @@ impl ComputeNode { self.prepare_pgdata(&compute_state, extension_server_port)?; let start_time = Utc::now(); - let pg = self.start_postgres(pspec.storage_auth_token.clone())?; + let pg = self.start_postgres(&pspec)?; let config_time = Utc::now(); if pspec.spec.mode == ComputeMode::Primary && !pspec.spec.skip_pg_catalog_updates { diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs index cba364c049..e329a2f23a 100644 --- a/control_plane/src/endpoint.rs +++ b/control_plane/src/endpoint.rs @@ -500,6 +500,7 @@ impl Endpoint { safekeeper_connstrings, storage_auth_token: auth_token.clone(), remote_extensions: None, + env_vars: None, }; let spec_path = self.endpoint_path().join("spec.json"); std::fs::write(spec_path, serde_json::to_string_pretty(&spec)?)?; diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs index b41ca8c9cf..4f4b693da4 100644 --- a/libs/compute_api/src/spec.rs +++ b/libs/compute_api/src/spec.rs @@ -58,7 +58,6 @@ pub struct ComputeSpec { pub pageserver_connstring: Option, #[serde(default)] pub safekeeper_connstrings: Vec, - #[serde(default)] pub mode: ComputeMode, @@ -68,6 +67,9 @@ pub struct ComputeSpec { // information about available remote extensions pub remote_extensions: Option, + + // Environment variables to set on the compute + pub env_vars: Option>>, } #[derive(Clone, Debug, Default, Deserialize, Serialize)] diff --git a/libs/compute_api/tests/cluster_spec.json b/libs/compute_api/tests/cluster_spec.json index 96db13a5da..d3772f22cd 100644 --- a/libs/compute_api/tests/cluster_spec.json +++ b/libs/compute_api/tests/cluster_spec.json @@ -243,5 +243,9 @@ "public_extensions": [ "postgis" ] - } + }, + "env_vars": { + "OPENAI_API_KEY" : "i hope heikki has a good day today", + "MY_NULL_ENV_VAR" : null + } }