From 1f8e8c50aeec61db795c3b983e654cf1d8eec01b Mon Sep 17 00:00:00 2001 From: Jarupat Jisarojito Date: Fri, 19 Jul 2024 11:06:06 -0700 Subject: [PATCH 1/3] Copy pg server cert and key to pgdata with correct permission ## Problem We need to copy certificate and key from secret mount directory to `pgdata` directory where `postgres` is the owner and we can set the key permission to 0600. ## Summary of changes - Added new pgparam `pg_compute_tls_settings` to specify where k8s secret for certificate and key are mounted. - Added a new field to `ComputeSpec` called `databricks_settings`. This is a struct that will be used to store any other settings that needs to be propagate to Compute but should not be persisted to `ComputeSpec` in the database. - Then when the compute container start up, as part of `prepare_pgdata` function, it will copied `server.key` and `server.crt` from k8s mounted directory to `pgdata` directory. Co-authored-by: Tristan Partin --- compute_tools/src/compute.rs | 8 ++++++++ control_plane/src/endpoint.rs | 1 + libs/compute_api/src/spec.rs | 4 ++++ 3 files changed, 13 insertions(+) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index a240e69df8..be734029df 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1462,6 +1462,14 @@ impl ComputeNode { // Update pg_hba.conf received with basebackup. update_pg_hba(pgdata_path, None)?; + if let Some(databricks_settings) = spec.databricks_settings.as_ref() { + copy_tls_certificates( + &databricks_settings.pg_compute_tls_settings.key_file, + &databricks_settings.pg_compute_tls_settings.cert_file, + pgdata_path, + )?; + } + // Place pg_dynshmem under /dev/shm. This allows us to use // 'dynamic_shared_memory_type = mmap' so that the files are placed in // /dev/shm, similar to how 'dynamic_shared_memory_type = posix' works. diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs index 6221d83c7f..1c7f489d68 100644 --- a/control_plane/src/endpoint.rs +++ b/control_plane/src/endpoint.rs @@ -793,6 +793,7 @@ impl Endpoint { autoprewarm: args.autoprewarm, offload_lfc_interval_seconds: args.offload_lfc_interval_seconds, suspend_timeout_seconds: -1, // Only used in neon_local. + databricks_settings: None, }; // this strange code is needed to support respec() in tests diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs index 061ac3e66d..474f523e73 100644 --- a/libs/compute_api/src/spec.rs +++ b/libs/compute_api/src/spec.rs @@ -193,6 +193,10 @@ pub struct ComputeSpec { /// /// We use this value to derive other values, such as the installed extensions metric. pub suspend_timeout_seconds: i64, + + // Databricks specific options for compute instance. + // These settings are not part of postgresql.conf. + pub databricks_settings: Option, } /// Feature flag to signal `compute_ctl` to enable certain experimental functionality. From 9fda72701845f692b1f73da7e41582cee8ca9e17 Mon Sep 17 00:00:00 2001 From: William Huang Date: Fri, 9 Aug 2024 14:25:47 -0700 Subject: [PATCH 2/3] [BRC-1425] Plumb through and set the requisite GUCs when starting the compute instance ## Problem We need the set the following Postgres GUCs to the correct value before starting Postgres in the compute instance: ``` databricks.workspace_url databricks.enable_databricks_identity_login databricks.enable_sql_restrictions ``` ## Summary of changes Plumbed through `workspace_url` and other GUC settings via `DatabricksSettings` in `ComputeSpec`. The spec is sent to the compute instance when it starts up and the GUCs are written to `postgresql.conf` before the postgres process is launched. Co-authored-by: Tristan Partin --- compute_tools/src/compute.rs | 73 +++++++++++++++++++++++++++++++++--- 1 file changed, 68 insertions(+), 5 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index be734029df..06a34ec41e 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -413,6 +413,52 @@ struct StartVmMonitorResult { vm_monitor: Option>>, } +/// Databricks-specific environment variables to be passed to the `postgres` sub-process. +pub struct DatabricksEnvVars { + /// The Databricks "endpoint ID" of the compute instance. Used by `postgres` to check + /// the token scopes of internal auth tokens. + pub endpoint_id: String, + /// Hostname of the Databricks workspace URL this compute instance belongs to. + /// Used by postgres to verify Databricks PAT tokens. + pub workspace_host: String, +} + +impl DatabricksEnvVars { + pub fn new(compute_spec: &ComputeSpec, compute_id: Option<&String>) -> Self { + // compute_id is a string format of "{endpoint_id}/{compute_idx}" + // endpoint_id is a uuid. We only need to pass down endpoint_id to postgres. + // Panics if compute_id is not set or not in the expected format. + let endpoint_id = compute_id.unwrap().split('/').next().unwrap().to_string(); + let workspace_host = compute_spec + .databricks_settings + .as_ref() + .map(|s| s.databricks_workspace_host.clone()) + .unwrap_or("".to_string()); + Self { + endpoint_id, + workspace_host, + } + } + + /// Constants for the names of Databricks-specific postgres environment variables. + const DATABRICKS_ENDPOINT_ID_ENVVAR: &'static str = "DATABRICKS_ENDPOINT_ID"; + const DATABRICKS_WORKSPACE_HOST_ENVVAR: &'static str = "DATABRICKS_WORKSPACE_HOST"; + + /// Convert DatabricksEnvVars to a list of string pairs that can be passed as env vars. Consumes `self`. + pub fn to_env_var_list(self) -> Vec<(String, String)> { + vec![ + ( + Self::DATABRICKS_ENDPOINT_ID_ENVVAR.to_string(), + self.endpoint_id.clone(), + ), + ( + Self::DATABRICKS_WORKSPACE_HOST_ENVVAR.to_string(), + self.workspace_host.clone(), + ), + ] + } +} + impl ComputeNode { pub fn new(params: ComputeNodeParams, config: ComputeConfig) -> Result { let connstr = params.connstr.as_str(); @@ -1581,14 +1627,31 @@ impl ComputeNode { pub fn start_postgres(&self, storage_auth_token: Option) -> Result { let pgdata_path = Path::new(&self.params.pgdata); + let env_vars: Vec<(String, String)> = if self.params.lakebase_mode { + let databricks_env_vars = { + let state = self.state.lock().unwrap(); + let spec = &state.pspec.as_ref().unwrap().spec; + DatabricksEnvVars::new(spec, Some(&self.params.compute_id)) + }; + + info!( + "Starting Postgres for databricks endpoint id: {}", + &databricks_env_vars.endpoint_id + ); + + let mut env_vars = databricks_env_vars.to_env_var_list(); + env_vars.extend(storage_auth_token.map(|t| ("NEON_AUTH_TOKEN".to_string(), t))); + env_vars + } else if let Some(storage_auth_token) = &storage_auth_token { + vec![("NEON_AUTH_TOKEN".to_owned(), storage_auth_token.to_owned())] + } else { + vec![] + }; + // Run postgres as a child process. let mut pg = maybe_cgexec(&self.params.pgbin) .args(["-D", &self.params.pgdata]) - .envs(if let Some(storage_auth_token) = &storage_auth_token { - vec![("NEON_AUTH_TOKEN", storage_auth_token)] - } else { - vec![] - }) + .envs(env_vars) .stderr(Stdio::piped()) .spawn() .expect("cannot start postgres process"); From a66df1f4bd88ab0182e55c4e99096b43a039f7db Mon Sep 17 00:00:00 2001 From: Jarupat Jisarojito Date: Tue, 23 Jul 2024 17:06:33 -0700 Subject: [PATCH 3/3] [BRC-1405] Mount databricks pg_hba and pg_ident from configmap to dblet pod ## Problem For certificate auth, we need to configure pg_hba and pg_ident for it to work. This PR https://github.com/databricks/universe/pull/655011 in universe will create a config map and deployed to `hadron-compute` namespace. HCC needs to mount this config map to all pg compute pod. ## Summary of changes Create `databricks_pg_hba` and `databricks_pg_ident` to configure where the files are located on the pod. These configs are pass down to `compute_ctl`. Compute_ctl uses these config to update `pg_hba.conf` and `pg_ident.conf` file. We append `include_if_exists {databricks_pg_hba}` to `pg_hba.conf` and similarly to `pg_ident.conf`. So that it will refer to databricks config file without much change to existing pg default config file. I renamed `secret_mounts` to `compute_mounts` because now it is used to configure secret and config map mounts. Co-authored-by: Tristan Partin --- compute_tools/src/compute.rs | 10 +++++++--- compute_tools/src/spec.rs | 3 ++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 06a34ec41e..b5415f8742 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1505,15 +1505,19 @@ impl ComputeNode { ) })?; - // Update pg_hba.conf received with basebackup. - update_pg_hba(pgdata_path, None)?; - if let Some(databricks_settings) = spec.databricks_settings.as_ref() { copy_tls_certificates( &databricks_settings.pg_compute_tls_settings.key_file, &databricks_settings.pg_compute_tls_settings.cert_file, pgdata_path, )?; + + // Update pg_hba.conf received with basebackup including additional databricks settings. + update_pg_hba(pgdata_path, Some(&databricks_settings.databricks_pg_hba))?; + update_pg_ident(pgdata_path, Some(&databricks_settings.databricks_pg_ident))?; + } else { + // Update pg_hba.conf received with basebackup. + update_pg_hba(pgdata_path, None)?; } // Place pg_dynshmem under /dev/shm. This allows us to use diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs index d00f86a2c0..751de76ca5 100644 --- a/compute_tools/src/spec.rs +++ b/compute_tools/src/spec.rs @@ -137,12 +137,13 @@ pub fn get_config_from_control_plane(base_uri: &str, compute_id: &str) -> Result /// Check `pg_hba.conf` and update if needed to allow external connections. pub fn update_pg_hba(pgdata_path: &Path, databricks_pg_hba: Option<&String>) -> Result<()> { // XXX: consider making it a part of config.json + info!("checking pg_hba.conf"); let pghba_path = pgdata_path.join("pg_hba.conf"); // Update pg_hba to contains databricks specfic settings before adding neon settings // PG uses the first record that matches to perform authentication, so we need to have // our rules before the default ones from neon. - // See https://www.postgresql.org/docs/16/auth-pg-hba-conf.html + // See https://www.postgresql.org/docs/current/auth-pg-hba-conf.html if let Some(databricks_pg_hba) = databricks_pg_hba { if config::line_in_file( &pghba_path,