LCOV - code coverage report
Current view: top level - compute_tools/src - pg_helpers.rs (source / functions) Coverage Total Hit
Test: 1e20c4f2b28aa592527961bb32170ebbd2c9172f.info Lines: 29.8 % 245 73
Test Date: 2025-07-16 12:29:03 Functions: 35.9 % 39 14

            Line data    Source code
       1              : use std::collections::HashMap;
       2              : use std::fmt::Write;
       3              : use std::fs;
       4              : use std::fs::File;
       5              : use std::io::{BufRead, BufReader};
       6              : use std::os::unix::fs::PermissionsExt;
       7              : use std::path::Path;
       8              : use std::process::Child;
       9              : use std::str::FromStr;
      10              : use std::time::{Duration, Instant};
      11              : 
      12              : use anyhow::{Result, bail};
      13              : use compute_api::responses::TlsConfig;
      14              : use compute_api::spec::{Database, GenericOption, GenericOptions, PgIdent, Role};
      15              : use futures::StreamExt;
      16              : use indexmap::IndexMap;
      17              : use ini::Ini;
      18              : use notify::{RecursiveMode, Watcher};
      19              : use postgres::config::Config;
      20              : use tokio::io::AsyncBufReadExt;
      21              : use tokio::task::JoinHandle;
      22              : use tokio::time::timeout;
      23              : use tokio_postgres;
      24              : use tokio_postgres::NoTls;
      25              : use tracing::{debug, error, info, instrument};
      26              : 
      27              : const POSTGRES_WAIT_TIMEOUT: Duration = Duration::from_millis(60 * 1000); // milliseconds
      28              : 
      29              : /// Escape a string for including it in a SQL literal.
      30              : ///
      31              : /// Wrapping the result with `E'{}'` or `'{}'` is not required,
      32              : /// as it returns a ready-to-use SQL string literal, e.g. `'db'''` or `E'db\\'`.
      33              : /// See <https://github.com/postgres/postgres/blob/da98d005cdbcd45af563d0c4ac86d0e9772cd15f/src/backend/utils/adt/quote.c#L47>
      34              : /// for the original implementation.
      35            6 : pub fn escape_literal(s: &str) -> String {
      36            6 :     let res = s.replace('\'', "''").replace('\\', "\\\\");
      37              : 
      38            6 :     if res.contains('\\') {
      39            2 :         format!("E'{res}'")
      40              :     } else {
      41            4 :         format!("'{res}'")
      42              :     }
      43            6 : }
      44              : 
      45              : /// Escape a string so that it can be used in postgresql.conf. Wrapping the result
      46              : /// with `'{}'` is not required, as it returns a ready-to-use config string.
      47            8 : pub fn escape_conf_value(s: &str) -> String {
      48            8 :     let res = s.replace('\'', "''").replace('\\', "\\\\");
      49            8 :     format!("'{res}'")
      50            8 : }
      51              : 
      52              : pub trait GenericOptionExt {
      53              :     fn to_pg_option(&self) -> String;
      54              :     fn to_pg_setting(&self) -> String;
      55              : }
      56              : 
      57              : impl GenericOptionExt for GenericOption {
      58              :     /// Represent `GenericOption` as SQL statement parameter.
      59            3 :     fn to_pg_option(&self) -> String {
      60            3 :         if let Some(val) = &self.value {
      61            3 :             match self.vartype.as_ref() {
      62            3 :                 "string" => format!("{} {}", self.name, escape_literal(val)),
      63            1 :                 _ => format!("{} {}", self.name, val),
      64              :             }
      65              :         } else {
      66            0 :             self.name.to_owned()
      67              :         }
      68            3 :     }
      69              : 
      70              :     /// Represent `GenericOption` as configuration option.
      71           25 :     fn to_pg_setting(&self) -> String {
      72           25 :         if let Some(val) = &self.value {
      73           25 :             match self.vartype.as_ref() {
      74           25 :                 "string" => format!("{} = {}", self.name, escape_conf_value(val)),
      75           17 :                 _ => format!("{} = {}", self.name, val),
      76              :             }
      77              :         } else {
      78            0 :             self.name.to_owned()
      79              :         }
      80           25 :     }
      81              : }
      82              : 
      83              : pub trait PgOptionsSerialize {
      84              :     fn as_pg_options(&self) -> String;
      85              :     fn as_pg_settings(&self) -> String;
      86              : }
      87              : 
      88              : impl PgOptionsSerialize for GenericOptions {
      89              :     /// Serialize an optional collection of `GenericOption`'s to
      90              :     /// Postgres SQL statement arguments.
      91            2 :     fn as_pg_options(&self) -> String {
      92            2 :         if let Some(ops) = &self {
      93            1 :             ops.iter()
      94            3 :                 .map(|op| op.to_pg_option())
      95            1 :                 .collect::<Vec<String>>()
      96            1 :                 .join(" ")
      97              :         } else {
      98            1 :             "".to_string()
      99              :         }
     100            2 :     }
     101              : 
     102              :     /// Serialize an optional collection of `GenericOption`'s to
     103              :     /// `postgresql.conf` compatible format.
     104            1 :     fn as_pg_settings(&self) -> String {
     105            1 :         if let Some(ops) = &self {
     106            1 :             ops.iter()
     107           25 :                 .map(|op| op.to_pg_setting())
     108            1 :                 .collect::<Vec<String>>()
     109            1 :                 .join("\n")
     110            1 :                 + "\n" // newline after last setting
     111              :         } else {
     112            0 :             "".to_string()
     113              :         }
     114            1 :     }
     115              : }
     116              : 
     117              : pub trait GenericOptionsSearch {
     118              :     fn find(&self, name: &str) -> Option<String>;
     119              :     fn find_ref(&self, name: &str) -> Option<&GenericOption>;
     120              : }
     121              : 
     122              : impl GenericOptionsSearch for GenericOptions {
     123              :     /// Lookup option by name
     124           15 :     fn find(&self, name: &str) -> Option<String> {
     125           15 :         let ops = self.as_ref()?;
     126          119 :         let op = ops.iter().find(|s| s.name == name)?;
     127            6 :         op.value.clone()
     128           15 :     }
     129              : 
     130              :     /// Lookup option by name, returning ref
     131            0 :     fn find_ref(&self, name: &str) -> Option<&GenericOption> {
     132            0 :         let ops = self.as_ref()?;
     133            0 :         ops.iter().find(|s| s.name == name)
     134            0 :     }
     135              : }
     136              : 
     137              : pub trait RoleExt {
     138              :     fn to_pg_options(&self) -> String;
     139              : }
     140              : 
     141              : impl RoleExt for Role {
     142              :     /// Serialize a list of role parameters into a Postgres-acceptable
     143              :     /// string of arguments.
     144            1 :     fn to_pg_options(&self) -> String {
     145              :         // XXX: consider putting LOGIN as a default option somewhere higher, e.g. in control-plane.
     146            1 :         let mut params: String = self.options.as_pg_options();
     147            1 :         params.push_str(" LOGIN");
     148              : 
     149            1 :         if let Some(pass) = &self.encrypted_password {
     150              :             // Some time ago we supported only md5 and treated all encrypted_password as md5.
     151              :             // Now we also support SCRAM-SHA-256 and to preserve compatibility
     152              :             // we treat all encrypted_password as md5 unless they starts with SCRAM-SHA-256.
     153            1 :             if pass.starts_with("SCRAM-SHA-256") {
     154            0 :                 write!(params, " PASSWORD '{pass}'")
     155            0 :                     .expect("String is documented to not to error during write operations");
     156            1 :             } else {
     157            1 :                 write!(params, " PASSWORD 'md5{pass}'")
     158            1 :                     .expect("String is documented to not to error during write operations");
     159            1 :             }
     160            0 :         } else {
     161            0 :             params.push_str(" PASSWORD NULL");
     162            0 :         }
     163              : 
     164            1 :         params
     165            1 :     }
     166              : }
     167              : 
     168              : pub trait DatabaseExt {
     169              :     fn to_pg_options(&self) -> String;
     170              : }
     171              : 
     172              : impl DatabaseExt for Database {
     173              :     /// Serialize a list of database parameters into a Postgres-acceptable
     174              :     /// string of arguments.
     175              :     /// NB: `TEMPLATE` is actually also an identifier, but so far we only need
     176              :     /// to use `template0` and `template1`, so it is not a problem. Yet in the future
     177              :     /// it may require a proper quoting too.
     178            1 :     fn to_pg_options(&self) -> String {
     179            1 :         let mut params: String = self.options.as_pg_options();
     180            1 :         write!(params, " OWNER {}", &self.owner.pg_quote())
     181            1 :             .expect("String is documented to not to error during write operations");
     182              : 
     183            1 :         params
     184            1 :     }
     185              : }
     186              : 
     187              : /// Generic trait used to provide quoting / encoding for strings used in the
     188              : /// Postgres SQL queries and DATABASE_URL.
     189              : pub trait Escaping {
     190              :     fn pg_quote(&self) -> String;
     191              :     fn pg_quote_dollar(&self) -> (String, String);
     192              : }
     193              : 
     194              : impl Escaping for PgIdent {
     195              :     /// This is intended to mimic Postgres quote_ident(), but for simplicity it
     196              :     /// always quotes provided string with `""` and escapes every `"`.
     197              :     /// **Not idempotent**, i.e. if string is already escaped it will be escaped again.
     198              :     /// N.B. it's not useful for escaping identifiers that are used inside WHERE
     199              :     /// clause, use `escape_literal()` instead.
     200            2 :     fn pg_quote(&self) -> String {
     201            2 :         format!("\"{}\"", self.replace('"', "\"\""))
     202            2 :     }
     203              : 
     204              :     /// This helper is intended to be used for dollar-escaping strings for usage
     205              :     /// inside PL/pgSQL procedures. In addition to dollar-escaping the string,
     206              :     /// it also returns a tag that is intended to be used inside the outer
     207              :     /// PL/pgSQL procedure. If you do not need an outer tag, just discard it.
     208              :     /// Here we somewhat mimic the logic of Postgres' `pg_get_functiondef()`,
     209              :     /// <https://github.com/postgres/postgres/blob/8b49392b270b4ac0b9f5c210e2a503546841e832/src/backend/utils/adt/ruleutils.c#L2924>
     210           14 :     fn pg_quote_dollar(&self) -> (String, String) {
     211           14 :         let mut tag: String = "x".to_string();
     212           14 :         let mut outer_tag = "xx".to_string();
     213              : 
     214              :         // Find the first suitable tag that is not present in the string.
     215              :         // Postgres' max role/DB name length is 63 bytes, so even in the
     216              :         // worst case it won't take long. Outer tag is always `tag + "x"`,
     217              :         // so if `tag` is not present in the string, `outer_tag` is not
     218              :         // present in the string either.
     219           27 :         while self.contains(&tag.to_string()) {
     220           13 :             tag += "x";
     221           13 :             outer_tag = tag.clone() + "x";
     222           13 :         }
     223              : 
     224           14 :         let escaped = format!("${tag}${self}${tag}$");
     225              : 
     226           14 :         (escaped, outer_tag)
     227           14 :     }
     228              : }
     229              : 
     230              : /// Build a list of existing Postgres roles
     231            0 : pub async fn get_existing_roles_async(client: &tokio_postgres::Client) -> Result<Vec<Role>> {
     232            0 :     let postgres_roles = client
     233            0 :         .query_raw::<str, &String, &[String; 0]>(
     234            0 :             "SELECT rolname, rolpassword FROM pg_catalog.pg_authid",
     235            0 :             &[],
     236            0 :         )
     237            0 :         .await?
     238            0 :         .filter_map(|row| async { row.ok() })
     239            0 :         .map(|row| Role {
     240            0 :             name: row.get("rolname"),
     241            0 :             encrypted_password: row.get("rolpassword"),
     242            0 :             options: None,
     243            0 :         })
     244            0 :         .collect()
     245            0 :         .await;
     246              : 
     247            0 :     Ok(postgres_roles)
     248            0 : }
     249              : 
     250              : /// Build a list of existing Postgres databases
     251            0 : pub async fn get_existing_dbs_async(
     252            0 :     client: &tokio_postgres::Client,
     253            0 : ) -> Result<HashMap<String, Database>> {
     254              :     // `pg_database.datconnlimit = -2` means that the database is in the
     255              :     // invalid state. See:
     256              :     //   https://github.com/postgres/postgres/commit/a4b4cc1d60f7e8ccfcc8ff8cb80c28ee411ad9a9
     257            0 :     let rowstream = client
     258            0 :         // We use a subquery instead of a fancy `datdba::regrole::text AS owner`,
     259            0 :         // because the latter automatically wraps the result in double quotes,
     260            0 :         // if the role name contains special characters.
     261            0 :         .query_raw::<str, &String, &[String; 0]>(
     262            0 :             "SELECT
     263            0 :                 datname AS name,
     264            0 :                 (SELECT rolname FROM pg_roles WHERE oid = datdba) AS owner,
     265            0 :                 NOT datallowconn AS restrict_conn,
     266            0 :                 datconnlimit = - 2 AS invalid
     267            0 :             FROM
     268            0 :                 pg_catalog.pg_database;",
     269            0 :             &[],
     270            0 :         )
     271            0 :         .await?;
     272              : 
     273            0 :     let dbs_map = rowstream
     274            0 :         .filter_map(|r| async { r.ok() })
     275            0 :         .map(|row| Database {
     276            0 :             name: row.get("name"),
     277            0 :             owner: row.get("owner"),
     278            0 :             restrict_conn: row.get("restrict_conn"),
     279            0 :             invalid: row.get("invalid"),
     280            0 :             options: None,
     281            0 :         })
     282            0 :         .map(|db| (db.name.clone(), db.clone()))
     283            0 :         .collect::<HashMap<_, _>>()
     284            0 :         .await;
     285              : 
     286            0 :     Ok(dbs_map)
     287            0 : }
     288              : 
     289              : /// Wait for Postgres to become ready to accept connections. It's ready to
     290              : /// accept connections when the state-field in `pgdata/postmaster.pid` says
     291              : /// 'ready'.
     292              : #[instrument(skip_all, fields(pgdata = %pgdata.display()))]
     293              : pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
     294              :     let pid_path = pgdata.join("postmaster.pid");
     295              : 
     296              :     // PostgreSQL writes line "ready" to the postmaster.pid file, when it has
     297              :     // completed initialization and is ready to accept connections. We want to
     298              :     // react quickly and perform the rest of our initialization as soon as
     299              :     // PostgreSQL starts accepting connections. Use 'notify' to be notified
     300              :     // whenever the PID file is changed, and whenever it changes, read it to
     301              :     // check if it's now "ready".
     302              :     //
     303              :     // You cannot actually watch a file before it exists, so we first watch the
     304              :     // data directory, and once the postmaster.pid file appears, we switch to
     305              :     // watch the file instead. We also wake up every 100 ms to poll, just in
     306              :     // case we miss some events for some reason. Not strictly necessary, but
     307              :     // better safe than sorry.
     308              :     let (tx, rx) = std::sync::mpsc::channel();
     309            0 :     let watcher_res = notify::recommended_watcher(move |res| {
     310            0 :         let _ = tx.send(res);
     311            0 :     });
     312              :     let (mut watcher, rx): (Box<dyn Watcher>, _) = match watcher_res {
     313              :         Ok(watcher) => (Box::new(watcher), rx),
     314              :         Err(e) => {
     315              :             match e.kind {
     316              :                 notify::ErrorKind::Io(os) if os.raw_os_error() == Some(38) => {
     317              :                     // docker on m1 macs does not support recommended_watcher
     318              :                     // but return "Function not implemented (os error 38)"
     319              :                     // see https://github.com/notify-rs/notify/issues/423
     320              :                     let (tx, rx) = std::sync::mpsc::channel();
     321              : 
     322              :                     // let's poll it faster than what we check the results for (100ms)
     323              :                     let config =
     324              :                         notify::Config::default().with_poll_interval(Duration::from_millis(50));
     325              : 
     326              :                     let watcher = notify::PollWatcher::new(
     327            0 :                         move |res| {
     328            0 :                             let _ = tx.send(res);
     329            0 :                         },
     330              :                         config,
     331              :                     )?;
     332              : 
     333              :                     (Box::new(watcher), rx)
     334              :                 }
     335              :                 _ => return Err(e.into()),
     336              :             }
     337              :         }
     338              :     };
     339              : 
     340              :     watcher.watch(pgdata, RecursiveMode::NonRecursive)?;
     341              : 
     342              :     let started_at = Instant::now();
     343              :     let mut postmaster_pid_seen = false;
     344              :     loop {
     345              :         if let Ok(Some(status)) = pg.try_wait() {
     346              :             // Postgres exited, that is not what we expected, bail out earlier.
     347              :             let code = status.code().unwrap_or(-1);
     348              :             bail!("Postgres exited unexpectedly with code {}", code);
     349              :         }
     350              : 
     351              :         let res = rx.recv_timeout(Duration::from_millis(100));
     352              :         debug!("woken up by notify: {res:?}");
     353              :         // If there are multiple events in the channel already, we only need to be
     354              :         // check once. Swallow the extra events before we go ahead to check the
     355              :         // pid file.
     356              :         while let Ok(res) = rx.try_recv() {
     357              :             debug!("swallowing extra event: {res:?}");
     358              :         }
     359              : 
     360              :         // Check that we can open pid file first.
     361              :         if let Ok(file) = File::open(&pid_path) {
     362              :             if !postmaster_pid_seen {
     363              :                 debug!("postmaster.pid appeared");
     364              :                 watcher
     365              :                     .unwatch(pgdata)
     366              :                     .expect("Failed to remove pgdata dir watch");
     367              :                 watcher
     368              :                     .watch(&pid_path, RecursiveMode::NonRecursive)
     369              :                     .expect("Failed to add postmaster.pid file watch");
     370              :                 postmaster_pid_seen = true;
     371              :             }
     372              : 
     373              :             let file = BufReader::new(file);
     374              :             let last_line = file.lines().last();
     375              : 
     376              :             // Pid file could be there and we could read it, but it could be empty, for example.
     377              :             if let Some(Ok(line)) = last_line {
     378              :                 let status = line.trim();
     379              :                 debug!("last line of postmaster.pid: {status:?}");
     380              : 
     381              :                 // Now Postgres is ready to accept connections
     382              :                 if status == "ready" {
     383              :                     break;
     384              :                 }
     385              :             }
     386              :         }
     387              : 
     388              :         // Give up after POSTGRES_WAIT_TIMEOUT.
     389              :         let duration = started_at.elapsed();
     390              :         if duration >= POSTGRES_WAIT_TIMEOUT {
     391              :             bail!("timed out while waiting for Postgres to start");
     392              :         }
     393              :     }
     394              : 
     395              :     tracing::info!("PostgreSQL is now running, continuing to configure it");
     396              : 
     397              :     Ok(())
     398              : }
     399              : 
     400              : /// Remove `pgdata` directory and create it again with right permissions.
     401            0 : pub fn create_pgdata(pgdata: &str) -> Result<()> {
     402              :     // Ignore removal error, likely it is a 'No such file or directory (os error 2)'.
     403              :     // If it is something different then create_dir() will error out anyway.
     404            0 :     let _ok = fs::remove_dir_all(pgdata);
     405            0 :     fs::create_dir(pgdata)?;
     406            0 :     fs::set_permissions(pgdata, fs::Permissions::from_mode(0o700))?;
     407              : 
     408            0 :     Ok(())
     409            0 : }
     410              : 
     411              : /// Update pgbouncer.ini with provided options
     412            0 : fn update_pgbouncer_ini(
     413            0 :     pgbouncer_config: IndexMap<String, String>,
     414            0 :     pgbouncer_ini_path: &str,
     415            0 : ) -> Result<()> {
     416            0 :     let mut conf = Ini::load_from_file(pgbouncer_ini_path)?;
     417            0 :     let section = conf.section_mut(Some("pgbouncer")).unwrap();
     418              : 
     419            0 :     for (option_name, value) in pgbouncer_config.iter() {
     420            0 :         section.insert(option_name, value);
     421            0 :         debug!(
     422            0 :             "Updating pgbouncer.ini with new values {}={}",
     423              :             option_name, value
     424              :         );
     425              :     }
     426              : 
     427            0 :     conf.write_to_file(pgbouncer_ini_path)?;
     428            0 :     Ok(())
     429            0 : }
     430              : 
     431              : /// Tune pgbouncer.
     432              : /// 1. Apply new config using pgbouncer admin console
     433              : /// 2. Add new values to pgbouncer.ini to preserve them after restart
     434            0 : pub async fn tune_pgbouncer(
     435            0 :     mut pgbouncer_config: IndexMap<String, String>,
     436            0 :     tls_config: Option<TlsConfig>,
     437            0 : ) -> Result<()> {
     438            0 :     let pgbouncer_connstr = if std::env::var_os("AUTOSCALING").is_some() {
     439              :         // for VMs use pgbouncer specific way to connect to
     440              :         // pgbouncer admin console without password
     441              :         // when pgbouncer is running under the same user.
     442            0 :         "host=/tmp port=6432 dbname=pgbouncer user=pgbouncer".to_string()
     443              :     } else {
     444              :         // for k8s use normal connection string with password
     445              :         // to connect to pgbouncer admin console
     446            0 :         let mut pgbouncer_connstr =
     447            0 :             "host=localhost port=6432 dbname=pgbouncer user=postgres sslmode=disable".to_string();
     448            0 :         if let Ok(pass) = std::env::var("PGBOUNCER_PASSWORD") {
     449            0 :             pgbouncer_connstr.push_str(format!(" password={pass}").as_str());
     450            0 :         }
     451            0 :         pgbouncer_connstr
     452              :     };
     453              : 
     454            0 :     info!(
     455            0 :         "Connecting to pgbouncer with connection string: {}",
     456              :         pgbouncer_connstr
     457              :     );
     458              : 
     459              :     // connect to pgbouncer, retrying several times
     460              :     // because pgbouncer may not be ready yet
     461            0 :     let mut retries = 3;
     462            0 :     let client = loop {
     463            0 :         match tokio_postgres::connect(&pgbouncer_connstr, NoTls).await {
     464            0 :             Ok((client, connection)) => {
     465            0 :                 tokio::spawn(async move {
     466            0 :                     if let Err(e) = connection.await {
     467            0 :                         eprintln!("connection error: {e}");
     468            0 :                     }
     469            0 :                 });
     470            0 :                 break client;
     471              :             }
     472            0 :             Err(e) => {
     473            0 :                 if retries == 0 {
     474            0 :                     return Err(e.into());
     475            0 :                 }
     476            0 :                 error!("Failed to connect to pgbouncer: pgbouncer_connstr {}", e);
     477            0 :                 retries -= 1;
     478            0 :                 tokio::time::sleep(Duration::from_secs(1)).await;
     479              :             }
     480              :         }
     481              :     };
     482              : 
     483            0 :     if let Some(tls_config) = tls_config {
     484              :         // pgbouncer starts in a half-ok state if it cannot find these files.
     485              :         // It will default to client_tls_sslmode=deny, which causes proxy to error.
     486              :         // There is a small window at startup where these files don't yet exist in the VM.
     487              :         // Best to wait until it exists.
     488              :         loop {
     489            0 :             if let Ok(true) = tokio::fs::try_exists(&tls_config.key_path).await {
     490            0 :                 break;
     491            0 :             }
     492            0 :             tokio::time::sleep(Duration::from_millis(500)).await
     493              :         }
     494              : 
     495            0 :         pgbouncer_config.insert("client_tls_cert_file".to_string(), tls_config.cert_path);
     496            0 :         pgbouncer_config.insert("client_tls_key_file".to_string(), tls_config.key_path);
     497            0 :         pgbouncer_config.insert("client_tls_sslmode".to_string(), "allow".to_string());
     498            0 :     }
     499              : 
     500              :     // save values to pgbouncer.ini
     501              :     // so that they are preserved after pgbouncer restart
     502            0 :     let pgbouncer_ini_path = if std::env::var_os("AUTOSCALING").is_some() {
     503              :         // in VMs we use /etc/pgbouncer.ini
     504            0 :         "/etc/pgbouncer.ini".to_string()
     505              :     } else {
     506              :         // in pods we use /var/db/postgres/pgbouncer/pgbouncer.ini
     507              :         // this is a shared volume between pgbouncer and postgres containers
     508              :         // FIXME: fix permissions for this file
     509            0 :         "/var/db/postgres/pgbouncer/pgbouncer.ini".to_string()
     510              :     };
     511            0 :     update_pgbouncer_ini(pgbouncer_config, &pgbouncer_ini_path)?;
     512              : 
     513            0 :     info!("Applying pgbouncer setting change");
     514              : 
     515            0 :     if let Err(err) = client.simple_query("RELOAD").await {
     516              :         // Don't fail on error, just print it into log
     517            0 :         error!("Failed to apply pgbouncer setting change,  {err}",);
     518            0 :     };
     519              : 
     520            0 :     Ok(())
     521            0 : }
     522              : 
     523              : /// Spawn a task that will read Postgres logs from `stderr`, join multiline logs
     524              : /// and send them to the logger. In the future we may also want to add context to
     525              : /// these logs.
     526            0 : pub fn handle_postgres_logs(stderr: std::process::ChildStderr) -> JoinHandle<Result<()>> {
     527            0 :     tokio::spawn(async move {
     528            0 :         let stderr = tokio::process::ChildStderr::from_std(stderr)?;
     529            0 :         handle_postgres_logs_async(stderr).await
     530            0 :     })
     531            0 : }
     532              : 
     533              : /// Read Postgres logs from `stderr` until EOF. Buffer is flushed on one of the following conditions:
     534              : /// - next line starts with timestamp
     535              : /// - EOF
     536              : /// - no new lines were written for the last 100 milliseconds
     537            0 : async fn handle_postgres_logs_async(stderr: tokio::process::ChildStderr) -> Result<()> {
     538            0 :     let mut lines = tokio::io::BufReader::new(stderr).lines();
     539            0 :     let timeout_duration = Duration::from_millis(100);
     540            0 :     let ts_regex =
     541            0 :         regex::Regex::new(r"^\d+-\d{2}-\d{2} \d{2}:\d{2}:\d{2}").expect("regex is valid");
     542              : 
     543            0 :     let mut buf = vec![];
     544              :     loop {
     545            0 :         let next_line = timeout(timeout_duration, lines.next_line()).await;
     546              : 
     547              :         // we should flush lines from the buffer if we cannot continue reading multiline message
     548            0 :         let should_flush_buf = match next_line {
     549              :             // Flushing if new line starts with timestamp
     550            0 :             Ok(Ok(Some(ref line))) => ts_regex.is_match(line),
     551              :             // Flushing on EOF, timeout or error
     552            0 :             _ => true,
     553              :         };
     554              : 
     555            0 :         if !buf.is_empty() && should_flush_buf {
     556              :             // join multiline message into a single line, separated by unicode Zero Width Space.
     557              :             // "PG:" suffix is used to distinguish postgres logs from other logs.
     558            0 :             let combined = format!("PG:{}\n", buf.join("\u{200B}"));
     559            0 :             buf.clear();
     560              : 
     561              :             // sync write to stderr to avoid interleaving with other logs
     562              :             use std::io::Write;
     563            0 :             let res = std::io::stderr().lock().write_all(combined.as_bytes());
     564            0 :             if let Err(e) = res {
     565            0 :                 tracing::error!("error while writing to stderr: {}", e);
     566            0 :             }
     567            0 :         }
     568              : 
     569              :         // if not timeout, append line to the buffer
     570            0 :         if next_line.is_ok() {
     571            0 :             match next_line?? {
     572            0 :                 Some(line) => buf.push(line),
     573              :                 // EOF
     574            0 :                 None => break,
     575              :             };
     576            0 :         }
     577              :     }
     578              : 
     579            0 :     Ok(())
     580            0 : }
     581              : 
     582              : /// `Postgres::config::Config` handles database names with whitespaces
     583              : /// and special characters properly.
     584            0 : pub fn postgres_conf_for_db(connstr: &url::Url, dbname: &str) -> Result<Config> {
     585            0 :     let mut conf = Config::from_str(connstr.as_str())?;
     586            0 :     conf.dbname(dbname);
     587            0 :     Ok(conf)
     588            0 : }
        

Generated by: LCOV version 2.1-beta