LCOV - code coverage report
Current view: top level - control_plane/src - postgresql_conf.rs (source / functions) Coverage Total Hit
Test: 75747cdbffeb0b6d2a2a311584368de68cd9aadc.info Lines: 50.4 % 139 70
Test Date: 2024-06-24 06:52:57 Functions: 23.5 % 17 4

            Line data    Source code
       1              : ///
       2              : /// Module for parsing postgresql.conf file.
       3              : ///
       4              : /// NOTE: This doesn't implement the full, correct postgresql.conf syntax. Just
       5              : /// enough to extract a few settings we need in Neon, assuming you don't do
       6              : /// funny stuff like include-directives or funny escaping.
       7              : use anyhow::{bail, Context, Result};
       8              : use once_cell::sync::Lazy;
       9              : use regex::Regex;
      10              : use std::collections::HashMap;
      11              : use std::fmt;
      12              : use std::io::BufRead;
      13              : use std::str::FromStr;
      14              : 
      15              : /// In-memory representation of a postgresql.conf file
      16              : #[derive(Default, Debug)]
      17              : pub struct PostgresConf {
      18              :     lines: Vec<String>,
      19              :     hash: HashMap<String, String>,
      20              : }
      21              : 
      22            0 : static CONF_LINE_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^((?:\w|\.)+)\s*=\s*(\S+)$").unwrap());
      23              : 
      24              : impl PostgresConf {
      25            0 :     pub fn new() -> PostgresConf {
      26            0 :         PostgresConf::default()
      27            0 :     }
      28              : 
      29              :     /// Read file into memory
      30            0 :     pub fn read(read: impl std::io::Read) -> Result<PostgresConf> {
      31            0 :         let mut result = Self::new();
      32              : 
      33            0 :         for line in std::io::BufReader::new(read).lines() {
      34            0 :             let line = line?;
      35              : 
      36              :             // Store each line in a vector, in original format
      37            0 :             result.lines.push(line.clone());
      38            0 : 
      39            0 :             // Also parse each line and insert key=value lines into a hash map.
      40            0 :             //
      41            0 :             // FIXME: This doesn't match exactly the flex/bison grammar in PostgreSQL.
      42            0 :             // But it's close enough for our usage.
      43            0 :             let line = line.trim();
      44            0 :             if line.starts_with('#') {
      45              :                 // comment, ignore
      46            0 :                 continue;
      47            0 :             } else if let Some(caps) = CONF_LINE_RE.captures(line) {
      48            0 :                 let name = caps.get(1).unwrap().as_str();
      49            0 :                 let raw_val = caps.get(2).unwrap().as_str();
      50              : 
      51            0 :                 if let Ok(val) = deescape_str(raw_val) {
      52            0 :                     // Note: if there's already an entry in the hash map for
      53            0 :                     // this key, this will replace it. That's the behavior what
      54            0 :                     // we want; when PostgreSQL reads the file, each line
      55            0 :                     // overrides any previous value for the same setting.
      56            0 :                     result.hash.insert(name.to_string(), val.to_string());
      57            0 :                 }
      58            0 :             }
      59              :         }
      60            0 :         Ok(result)
      61            0 :     }
      62              : 
      63              :     /// Return the current value of 'option'
      64            0 :     pub fn get(&self, option: &str) -> Option<&str> {
      65            0 :         self.hash.get(option).map(|x| x.as_ref())
      66            0 :     }
      67              : 
      68              :     /// Return the current value of a field, parsed to the right datatype.
      69              :     ///
      70              :     /// This calls the FromStr::parse() function on the value of the field. If
      71              :     /// the field does not exist, or parsing fails, returns an error.
      72              :     ///
      73            0 :     pub fn parse_field<T>(&self, field_name: &str, context: &str) -> Result<T>
      74            0 :     where
      75            0 :         T: FromStr,
      76            0 :         <T as FromStr>::Err: std::error::Error + Send + Sync + 'static,
      77            0 :     {
      78            0 :         self.get(field_name)
      79            0 :             .with_context(|| format!("could not find '{}' option {}", field_name, context))?
      80            0 :             .parse::<T>()
      81            0 :             .with_context(|| format!("could not parse '{}' option {}", field_name, context))
      82            0 :     }
      83              : 
      84            0 :     pub fn parse_field_optional<T>(&self, field_name: &str, context: &str) -> Result<Option<T>>
      85            0 :     where
      86            0 :         T: FromStr,
      87            0 :         <T as FromStr>::Err: std::error::Error + Send + Sync + 'static,
      88            0 :     {
      89            0 :         if let Some(val) = self.get(field_name) {
      90            0 :             let result = val
      91            0 :                 .parse::<T>()
      92            0 :                 .with_context(|| format!("could not parse '{}' option {}", field_name, context))?;
      93              : 
      94            0 :             Ok(Some(result))
      95              :         } else {
      96            0 :             Ok(None)
      97              :         }
      98            0 :     }
      99              : 
     100              :     ///
     101              :     /// Note: if you call this multiple times for the same option, the config
     102              :     /// file will a line for each call. It would be nice to have a function
     103              :     /// to change an existing line, but that's a TODO.
     104              :     ///
     105            0 :     pub fn append(&mut self, option: &str, value: &str) {
     106            0 :         self.lines
     107            0 :             .push(format!("{}={}\n", option, escape_str(value)));
     108            0 :         self.hash.insert(option.to_string(), value.to_string());
     109            0 :     }
     110              : 
     111              :     /// Append an arbitrary non-setting line to the config file
     112            0 :     pub fn append_line(&mut self, line: &str) {
     113            0 :         self.lines.push(line.to_string());
     114            0 :     }
     115              : }
     116              : 
     117              : impl fmt::Display for PostgresConf {
     118              :     /// Return the whole configuration file as a string
     119            0 :     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
     120            0 :         for line in self.lines.iter() {
     121            0 :             f.write_str(line)?;
     122              :         }
     123            0 :         Ok(())
     124            0 :     }
     125              : }
     126              : 
     127              : /// Escape a value for putting in postgresql.conf.
     128           32 : fn escape_str(s: &str) -> String {
     129           32 :     // If the string doesn't contain anything that needs quoting or escaping, return it
     130           32 :     // as it is.
     131           32 :     //
     132           32 :     // The first part of the regex, before the '|', matches the INTEGER rule in the
     133           32 :     // PostgreSQL flex grammar (guc-file.l). It matches plain integers like "123" and
     134           32 :     // "-123", and also accepts units like "10MB". The second part of the regex matches
     135           32 :     // the UNQUOTED_STRING rule, and accepts strings that contain a single word, beginning
     136           32 :     // with a letter. That covers words like "off" or "posix". Everything else is quoted.
     137           32 :     //
     138           32 :     // This regex is a bit more conservative than the rules in guc-file.l, so we quote some
     139           32 :     // strings that PostgreSQL would accept without quoting, but that's OK.
     140           32 : 
     141           32 :     static UNQUOTED_RE: Lazy<Regex> =
     142           32 :         Lazy::new(|| Regex::new(r"(^[-+]?[0-9]+[a-zA-Z]*$)|(^[a-zA-Z][a-zA-Z0-9]*$)").unwrap());
     143           32 : 
     144           32 :     if UNQUOTED_RE.is_match(s) {
     145           18 :         s.to_string()
     146              :     } else {
     147              :         // Otherwise escape and quote it
     148           14 :         let s = s
     149           14 :             .replace('\\', "\\\\")
     150           14 :             .replace('\n', "\\n")
     151           14 :             .replace('\'', "''");
     152           14 : 
     153           14 :         "\'".to_owned() + &s + "\'"
     154              :     }
     155           32 : }
     156              : 
     157              : /// De-escape a possibly-quoted value.
     158              : ///
     159              : /// See `DeescapeQuotedString` function in PostgreSQL sources for how PostgreSQL
     160              : /// does this.
     161            8 : fn deescape_str(s: &str) -> Result<String> {
     162            8 :     // If the string has a quote at the beginning and end, strip them out.
     163            8 :     if s.len() >= 2 && s.starts_with('\'') && s.ends_with('\'') {
     164            6 :         let mut result = String::new();
     165            6 : 
     166            6 :         let mut iter = s[1..(s.len() - 1)].chars().peekable();
     167           40 :         while let Some(c) = iter.next() {
     168           36 :             let newc = if c == '\\' {
     169           16 :                 match iter.next() {
     170            2 :                     Some('b') => '\x08',
     171            2 :                     Some('f') => '\x0c',
     172            4 :                     Some('n') => '\n',
     173            2 :                     Some('r') => '\r',
     174            2 :                     Some('t') => '\t',
     175            4 :                     Some('0'..='7') => {
     176              :                         // TODO
     177            2 :                         bail!("octal escapes not supported");
     178              :                     }
     179            2 :                     Some(n) => n,
     180            0 :                     None => break,
     181              :                 }
     182           20 :             } else if c == '\'' && iter.peek() == Some(&'\'') {
     183              :                 // doubled quote becomes just one quote
     184            2 :                 iter.next().unwrap()
     185              :             } else {
     186           18 :                 c
     187              :             };
     188              : 
     189           34 :             result.push(newc);
     190              :         }
     191            4 :         Ok(result)
     192              :     } else {
     193            2 :         Ok(s.to_string())
     194              :     }
     195            8 : }
     196              : 
     197              : #[test]
     198            2 : fn test_postgresql_conf_escapes() -> Result<()> {
     199            2 :     assert_eq!(escape_str("foo bar"), "'foo bar'");
     200              :     // these don't need to be quoted
     201            2 :     assert_eq!(escape_str("foo"), "foo");
     202            2 :     assert_eq!(escape_str("123"), "123");
     203            2 :     assert_eq!(escape_str("+123"), "+123");
     204            2 :     assert_eq!(escape_str("-10"), "-10");
     205            2 :     assert_eq!(escape_str("1foo"), "1foo");
     206            2 :     assert_eq!(escape_str("foo1"), "foo1");
     207            2 :     assert_eq!(escape_str("10MB"), "10MB");
     208            2 :     assert_eq!(escape_str("-10kB"), "-10kB");
     209              : 
     210              :     // these need quoting and/or escaping
     211            2 :     assert_eq!(escape_str("foo bar"), "'foo bar'");
     212            2 :     assert_eq!(escape_str("fo'o"), "'fo''o'");
     213            2 :     assert_eq!(escape_str("fo\no"), "'fo\\no'");
     214            2 :     assert_eq!(escape_str("fo\\o"), "'fo\\\\o'");
     215            2 :     assert_eq!(escape_str("10 cats"), "'10 cats'");
     216              : 
     217              :     // Test de-escaping
     218            2 :     assert_eq!(deescape_str(&escape_str("foo"))?, "foo");
     219            2 :     assert_eq!(deescape_str(&escape_str("fo'o\nba\\r"))?, "fo'o\nba\\r");
     220            2 :     assert_eq!(deescape_str("'\\b\\f\\n\\r\\t'")?, "\x08\x0c\n\r\t");
     221              : 
     222              :     // octal-escapes are currently not supported
     223            2 :     assert!(deescape_str("'foo\\7\\07\\007'").is_err());
     224              : 
     225            2 :     Ok(())
     226            2 : }
        

Generated by: LCOV version 2.1-beta