LCOV - code coverage report
Current view: top level - compute_tools/src - pg_helpers.rs (source / functions) Coverage Total Hit
Test: ae4948feae6a1d420c855050eb8c189119446a71.info Lines: 30.6 % 252 77
Test Date: 2025-03-18 18:33:46 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            6 : 
      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           23 :     fn to_pg_setting(&self) -> String {
      72           23 :         if let Some(val) = &self.value {
      73           23 :             match self.vartype.as_ref() {
      74           23 :                 "string" => format!("{} = {}", self.name, escape_conf_value(val)),
      75           15 :                 _ => format!("{} = {}", self.name, val),
      76              :             }
      77              :         } else {
      78            0 :             self.name.to_owned()
      79              :         }
      80           23 :     }
      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           23 :                 .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            9 :     fn find(&self, name: &str) -> Option<String> {
     125            9 :         let ops = self.as_ref()?;
     126            6 :         let op = ops.iter().find(|s| s.name == name)?;
     127            2 :         op.value.clone()
     128            9 :     }
     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            1 :         // 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            1 : 
     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            6 :     fn pg_quote_dollar(&self) -> (String, String) {
     211            6 :         let mut tag: String = "x".to_string();
     212            6 :         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.
     217            7 :         while self.contains(&format!("${tag}$")) || self.contains(&format!("${outer_tag}$")) {
     218            1 :             tag += "x";
     219            1 :             outer_tag = tag.clone() + "x";
     220            1 :         }
     221              : 
     222            6 :         let escaped = format!("${tag}${self}${tag}$");
     223            6 : 
     224            6 :         (escaped, outer_tag)
     225            6 :     }
     226              : }
     227              : 
     228              : /// Build a list of existing Postgres roles
     229            0 : pub async fn get_existing_roles_async(client: &tokio_postgres::Client) -> Result<Vec<Role>> {
     230            0 :     let postgres_roles = client
     231            0 :         .query_raw::<str, &String, &[String; 0]>(
     232            0 :             "SELECT rolname, rolpassword FROM pg_catalog.pg_authid",
     233            0 :             &[],
     234            0 :         )
     235            0 :         .await?
     236            0 :         .filter_map(|row| async { row.ok() })
     237            0 :         .map(|row| Role {
     238            0 :             name: row.get("rolname"),
     239            0 :             encrypted_password: row.get("rolpassword"),
     240            0 :             options: None,
     241            0 :         })
     242            0 :         .collect()
     243            0 :         .await;
     244              : 
     245            0 :     Ok(postgres_roles)
     246            0 : }
     247              : 
     248              : /// Build a list of existing Postgres databases
     249            0 : pub async fn get_existing_dbs_async(
     250            0 :     client: &tokio_postgres::Client,
     251            0 : ) -> Result<HashMap<String, Database>> {
     252              :     // `pg_database.datconnlimit = -2` means that the database is in the
     253              :     // invalid state. See:
     254              :     //   https://github.com/postgres/postgres/commit/a4b4cc1d60f7e8ccfcc8ff8cb80c28ee411ad9a9
     255            0 :     let rowstream = client
     256            0 :         // We use a subquery instead of a fancy `datdba::regrole::text AS owner`,
     257            0 :         // because the latter automatically wraps the result in double quotes,
     258            0 :         // if the role name contains special characters.
     259            0 :         .query_raw::<str, &String, &[String; 0]>(
     260            0 :             "SELECT
     261            0 :                 datname AS name,
     262            0 :                 (SELECT rolname FROM pg_roles WHERE oid = datdba) AS owner,
     263            0 :                 NOT datallowconn AS restrict_conn,
     264            0 :                 datconnlimit = - 2 AS invalid
     265            0 :             FROM
     266            0 :                 pg_catalog.pg_database;",
     267            0 :             &[],
     268            0 :         )
     269            0 :         .await?;
     270              : 
     271            0 :     let dbs_map = rowstream
     272            0 :         .filter_map(|r| async { r.ok() })
     273            0 :         .map(|row| Database {
     274            0 :             name: row.get("name"),
     275            0 :             owner: row.get("owner"),
     276            0 :             restrict_conn: row.get("restrict_conn"),
     277            0 :             invalid: row.get("invalid"),
     278            0 :             options: None,
     279            0 :         })
     280            0 :         .map(|db| (db.name.clone(), db.clone()))
     281            0 :         .collect::<HashMap<_, _>>()
     282            0 :         .await;
     283              : 
     284            0 :     Ok(dbs_map)
     285            0 : }
     286              : 
     287              : /// Wait for Postgres to become ready to accept connections. It's ready to
     288              : /// accept connections when the state-field in `pgdata/postmaster.pid` says
     289              : /// 'ready'.
     290              : #[instrument(skip_all, fields(pgdata = %pgdata.display()))]
     291              : pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
     292              :     let pid_path = pgdata.join("postmaster.pid");
     293              : 
     294              :     // PostgreSQL writes line "ready" to the postmaster.pid file, when it has
     295              :     // completed initialization and is ready to accept connections. We want to
     296              :     // react quickly and perform the rest of our initialization as soon as
     297              :     // PostgreSQL starts accepting connections. Use 'notify' to be notified
     298              :     // whenever the PID file is changed, and whenever it changes, read it to
     299              :     // check if it's now "ready".
     300              :     //
     301              :     // You cannot actually watch a file before it exists, so we first watch the
     302              :     // data directory, and once the postmaster.pid file appears, we switch to
     303              :     // watch the file instead. We also wake up every 100 ms to poll, just in
     304              :     // case we miss some events for some reason. Not strictly necessary, but
     305              :     // better safe than sorry.
     306              :     let (tx, rx) = std::sync::mpsc::channel();
     307            0 :     let watcher_res = notify::recommended_watcher(move |res| {
     308            0 :         let _ = tx.send(res);
     309            0 :     });
     310              :     let (mut watcher, rx): (Box<dyn Watcher>, _) = match watcher_res {
     311              :         Ok(watcher) => (Box::new(watcher), rx),
     312              :         Err(e) => {
     313              :             match e.kind {
     314              :                 notify::ErrorKind::Io(os) if os.raw_os_error() == Some(38) => {
     315              :                     // docker on m1 macs does not support recommended_watcher
     316              :                     // but return "Function not implemented (os error 38)"
     317              :                     // see https://github.com/notify-rs/notify/issues/423
     318              :                     let (tx, rx) = std::sync::mpsc::channel();
     319              : 
     320              :                     // let's poll it faster than what we check the results for (100ms)
     321              :                     let config =
     322              :                         notify::Config::default().with_poll_interval(Duration::from_millis(50));
     323              : 
     324              :                     let watcher = notify::PollWatcher::new(
     325            0 :                         move |res| {
     326            0 :                             let _ = tx.send(res);
     327            0 :                         },
     328              :                         config,
     329              :                     )?;
     330              : 
     331              :                     (Box::new(watcher), rx)
     332              :                 }
     333              :                 _ => return Err(e.into()),
     334              :             }
     335              :         }
     336              :     };
     337              : 
     338              :     watcher.watch(pgdata, RecursiveMode::NonRecursive)?;
     339              : 
     340              :     let started_at = Instant::now();
     341              :     let mut postmaster_pid_seen = false;
     342              :     loop {
     343              :         if let Ok(Some(status)) = pg.try_wait() {
     344              :             // Postgres exited, that is not what we expected, bail out earlier.
     345              :             let code = status.code().unwrap_or(-1);
     346              :             bail!("Postgres exited unexpectedly with code {}", code);
     347              :         }
     348              : 
     349              :         let res = rx.recv_timeout(Duration::from_millis(100));
     350              :         debug!("woken up by notify: {res:?}");
     351              :         // If there are multiple events in the channel already, we only need to be
     352              :         // check once. Swallow the extra events before we go ahead to check the
     353              :         // pid file.
     354              :         while let Ok(res) = rx.try_recv() {
     355              :             debug!("swallowing extra event: {res:?}");
     356              :         }
     357              : 
     358              :         // Check that we can open pid file first.
     359              :         if let Ok(file) = File::open(&pid_path) {
     360              :             if !postmaster_pid_seen {
     361              :                 debug!("postmaster.pid appeared");
     362              :                 watcher
     363              :                     .unwatch(pgdata)
     364              :                     .expect("Failed to remove pgdata dir watch");
     365              :                 watcher
     366              :                     .watch(&pid_path, RecursiveMode::NonRecursive)
     367              :                     .expect("Failed to add postmaster.pid file watch");
     368              :                 postmaster_pid_seen = true;
     369              :             }
     370              : 
     371              :             let file = BufReader::new(file);
     372              :             let last_line = file.lines().last();
     373              : 
     374              :             // Pid file could be there and we could read it, but it could be empty, for example.
     375              :             if let Some(Ok(line)) = last_line {
     376              :                 let status = line.trim();
     377              :                 debug!("last line of postmaster.pid: {status:?}");
     378              : 
     379              :                 // Now Postgres is ready to accept connections
     380              :                 if status == "ready" {
     381              :                     break;
     382              :                 }
     383              :             }
     384              :         }
     385              : 
     386              :         // Give up after POSTGRES_WAIT_TIMEOUT.
     387              :         let duration = started_at.elapsed();
     388              :         if duration >= POSTGRES_WAIT_TIMEOUT {
     389              :             bail!("timed out while waiting for Postgres to start");
     390              :         }
     391              :     }
     392              : 
     393              :     tracing::info!("PostgreSQL is now running, continuing to configure it");
     394              : 
     395              :     Ok(())
     396              : }
     397              : 
     398              : /// Remove `pgdata` directory and create it again with right permissions.
     399            0 : pub fn create_pgdata(pgdata: &str) -> Result<()> {
     400            0 :     // Ignore removal error, likely it is a 'No such file or directory (os error 2)'.
     401            0 :     // If it is something different then create_dir() will error out anyway.
     402            0 :     let _ok = fs::remove_dir_all(pgdata);
     403            0 :     fs::create_dir(pgdata)?;
     404            0 :     fs::set_permissions(pgdata, fs::Permissions::from_mode(0o700))?;
     405              : 
     406            0 :     Ok(())
     407            0 : }
     408              : 
     409              : /// Update pgbouncer.ini with provided options
     410            0 : fn update_pgbouncer_ini(
     411            0 :     pgbouncer_config: IndexMap<String, String>,
     412            0 :     pgbouncer_ini_path: &str,
     413            0 : ) -> Result<()> {
     414            0 :     let mut conf = Ini::load_from_file(pgbouncer_ini_path)?;
     415            0 :     let section = conf.section_mut(Some("pgbouncer")).unwrap();
     416              : 
     417            0 :     for (option_name, value) in pgbouncer_config.iter() {
     418            0 :         section.insert(option_name, value);
     419            0 :         debug!(
     420            0 :             "Updating pgbouncer.ini with new values {}={}",
     421              :             option_name, value
     422              :         );
     423              :     }
     424              : 
     425            0 :     conf.write_to_file(pgbouncer_ini_path)?;
     426            0 :     Ok(())
     427            0 : }
     428              : 
     429              : /// Tune pgbouncer.
     430              : /// 1. Apply new config using pgbouncer admin console
     431              : /// 2. Add new values to pgbouncer.ini to preserve them after restart
     432            0 : pub async fn tune_pgbouncer(
     433            0 :     mut pgbouncer_config: IndexMap<String, String>,
     434            0 :     tls_config: Option<TlsConfig>,
     435            0 : ) -> Result<()> {
     436            0 :     let pgbouncer_connstr = if std::env::var_os("AUTOSCALING").is_some() {
     437              :         // for VMs use pgbouncer specific way to connect to
     438              :         // pgbouncer admin console without password
     439              :         // when pgbouncer is running under the same user.
     440            0 :         "host=/tmp port=6432 dbname=pgbouncer user=pgbouncer".to_string()
     441              :     } else {
     442              :         // for k8s use normal connection string with password
     443              :         // to connect to pgbouncer admin console
     444            0 :         let mut pgbouncer_connstr =
     445            0 :             "host=localhost port=6432 dbname=pgbouncer user=postgres sslmode=disable".to_string();
     446            0 :         if let Ok(pass) = std::env::var("PGBOUNCER_PASSWORD") {
     447            0 :             pgbouncer_connstr.push_str(format!(" password={}", pass).as_str());
     448            0 :         }
     449            0 :         pgbouncer_connstr
     450              :     };
     451              : 
     452            0 :     info!(
     453            0 :         "Connecting to pgbouncer with connection string: {}",
     454              :         pgbouncer_connstr
     455              :     );
     456              : 
     457              :     // connect to pgbouncer, retrying several times
     458              :     // because pgbouncer may not be ready yet
     459            0 :     let mut retries = 3;
     460            0 :     let client = loop {
     461            0 :         match tokio_postgres::connect(&pgbouncer_connstr, NoTls).await {
     462            0 :             Ok((client, connection)) => {
     463            0 :                 tokio::spawn(async move {
     464            0 :                     if let Err(e) = connection.await {
     465            0 :                         eprintln!("connection error: {}", e);
     466            0 :                     }
     467            0 :                 });
     468            0 :                 break client;
     469              :             }
     470            0 :             Err(e) => {
     471            0 :                 if retries == 0 {
     472            0 :                     return Err(e.into());
     473            0 :                 }
     474            0 :                 error!("Failed to connect to pgbouncer: pgbouncer_connstr {}", e);
     475            0 :                 retries -= 1;
     476            0 :                 tokio::time::sleep(Duration::from_secs(1)).await;
     477              :             }
     478              :         }
     479              :     };
     480              : 
     481            0 :     if let Some(tls_config) = tls_config {
     482              :         // pgbouncer starts in a half-ok state if it cannot find these files.
     483              :         // It will default to client_tls_sslmode=deny, which causes proxy to error.
     484              :         // There is a small window at startup where these files don't yet exist in the VM.
     485              :         // Best to wait until it exists.
     486              :         loop {
     487            0 :             if let Ok(true) = tokio::fs::try_exists(&tls_config.key_path).await {
     488            0 :                 break;
     489            0 :             }
     490            0 :             tokio::time::sleep(Duration::from_millis(500)).await
     491              :         }
     492              : 
     493            0 :         pgbouncer_config.insert("client_tls_cert_file".to_string(), tls_config.cert_path);
     494            0 :         pgbouncer_config.insert("client_tls_key_file".to_string(), tls_config.key_path);
     495            0 :         pgbouncer_config.insert("client_tls_sslmode".to_string(), "allow".to_string());
     496            0 :     }
     497              : 
     498              :     // save values to pgbouncer.ini
     499              :     // so that they are preserved after pgbouncer restart
     500            0 :     let pgbouncer_ini_path = if std::env::var_os("AUTOSCALING").is_some() {
     501              :         // in VMs we use /etc/pgbouncer.ini
     502            0 :         "/etc/pgbouncer.ini".to_string()
     503              :     } else {
     504              :         // in pods we use /var/db/postgres/pgbouncer/pgbouncer.ini
     505              :         // this is a shared volume between pgbouncer and postgres containers
     506              :         // FIXME: fix permissions for this file
     507            0 :         "/var/db/postgres/pgbouncer/pgbouncer.ini".to_string()
     508              :     };
     509            0 :     update_pgbouncer_ini(pgbouncer_config, &pgbouncer_ini_path)?;
     510              : 
     511            0 :     info!("Applying pgbouncer setting change");
     512              : 
     513            0 :     if let Err(err) = client.simple_query("RELOAD").await {
     514              :         // Don't fail on error, just print it into log
     515            0 :         error!("Failed to apply pgbouncer setting change,  {err}",);
     516            0 :     };
     517              : 
     518            0 :     Ok(())
     519            0 : }
     520              : 
     521              : /// Spawn a task that will read Postgres logs from `stderr`, join multiline logs
     522              : /// and send them to the logger. In the future we may also want to add context to
     523              : /// these logs.
     524            0 : pub fn handle_postgres_logs(stderr: std::process::ChildStderr) -> JoinHandle<Result<()>> {
     525            0 :     tokio::spawn(async move {
     526            0 :         let stderr = tokio::process::ChildStderr::from_std(stderr)?;
     527            0 :         handle_postgres_logs_async(stderr).await
     528            0 :     })
     529            0 : }
     530              : 
     531              : /// Read Postgres logs from `stderr` until EOF. Buffer is flushed on one of the following conditions:
     532              : /// - next line starts with timestamp
     533              : /// - EOF
     534              : /// - no new lines were written for the last 100 milliseconds
     535            0 : async fn handle_postgres_logs_async(stderr: tokio::process::ChildStderr) -> Result<()> {
     536            0 :     let mut lines = tokio::io::BufReader::new(stderr).lines();
     537            0 :     let timeout_duration = Duration::from_millis(100);
     538            0 :     let ts_regex =
     539            0 :         regex::Regex::new(r"^\d+-\d{2}-\d{2} \d{2}:\d{2}:\d{2}").expect("regex is valid");
     540            0 : 
     541            0 :     let mut buf = vec![];
     542              :     loop {
     543            0 :         let next_line = timeout(timeout_duration, lines.next_line()).await;
     544              : 
     545              :         // we should flush lines from the buffer if we cannot continue reading multiline message
     546            0 :         let should_flush_buf = match next_line {
     547              :             // Flushing if new line starts with timestamp
     548            0 :             Ok(Ok(Some(ref line))) => ts_regex.is_match(line),
     549              :             // Flushing on EOF, timeout or error
     550            0 :             _ => true,
     551              :         };
     552              : 
     553            0 :         if !buf.is_empty() && should_flush_buf {
     554              :             // join multiline message into a single line, separated by unicode Zero Width Space.
     555              :             // "PG:" suffix is used to distinguish postgres logs from other logs.
     556            0 :             let combined = format!("PG:{}\n", buf.join("\u{200B}"));
     557            0 :             buf.clear();
     558              : 
     559              :             // sync write to stderr to avoid interleaving with other logs
     560              :             use std::io::Write;
     561            0 :             let res = std::io::stderr().lock().write_all(combined.as_bytes());
     562            0 :             if let Err(e) = res {
     563            0 :                 tracing::error!("error while writing to stderr: {}", e);
     564            0 :             }
     565            0 :         }
     566              : 
     567              :         // if not timeout, append line to the buffer
     568            0 :         if next_line.is_ok() {
     569            0 :             match next_line?? {
     570            0 :                 Some(line) => buf.push(line),
     571              :                 // EOF
     572            0 :                 None => break,
     573              :             };
     574            0 :         }
     575              :     }
     576              : 
     577            0 :     Ok(())
     578            0 : }
     579              : 
     580              : /// `Postgres::config::Config` handles database names with whitespaces
     581              : /// and special characters properly.
     582            0 : pub fn postgres_conf_for_db(connstr: &url::Url, dbname: &str) -> Result<Config> {
     583            0 :     let mut conf = Config::from_str(connstr.as_str())?;
     584            0 :     conf.dbname(dbname);
     585            0 :     Ok(conf)
     586            0 : }
        

Generated by: LCOV version 2.1-beta